hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

partial fix to loading checkpoints

Browse Source
This commit is contained in:
Nicklas Hansen
2025-01-21 00:10:53 -08:00
parent ae4238946f
commit dddc226d25
7 changed files with 23 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tdmpc2/common/buffer.py
View File

@@ -12,7 +12,7 @@ class Buffer():
def __init__(self, cfg): def __init__(self, cfg):
self.cfg = cfg self.cfg = cfg
self._device = torch.get_default_device() self._device = torch.device('cuda:0')
self._capacity = min(cfg.buffer_size, cfg.steps) self._capacity = min(cfg.buffer_size, cfg.steps)
self._sampler = SliceSampler( self._sampler = SliceSampler(
num_slices=self.cfg.batch_size, num_slices=self.cfg.batch_size,
@@ -59,7 +59,7 @@ class Buffer():
total_bytes = bytes_per_step*self._capacity total_bytes = bytes_per_step*self._capacity
print(f'Storage required: {total_bytes/1e9:.2f} GB') print(f'Storage required: {total_bytes/1e9:.2f} GB')
# Heuristic: decide whether to use CUDA or CPU memory # Heuristic: decide whether to use CUDA or CPU memory
storage_device = torch.get_default_device() if 2.5*total_bytes < mem_free else 'cpu' storage_device = 'cuda:0' if 2.5*total_bytes < mem_free else 'cpu'
print(f'Using {storage_device.upper()} memory for storage.') print(f'Using {storage_device.upper()} memory for storage.')
self._storage_device = torch.device(storage_device) self._storage_device = torch.device(storage_device)
return self._reserve_buffer( return self._reserve_buffer(

5
tdmpc2/common/scale.py
View File

@@ -1,14 +1,15 @@
import torch import torch
from torch.nn import Buffer from torch.nn import Buffer
class RunningScale(torch.nn.Module): class RunningScale(torch.nn.Module):
"""Running trimmed scale estimator.""" """Running trimmed scale estimator."""
def __init__(self, cfg): def __init__(self, cfg):
super().__init__() super().__init__()
self.cfg = cfg self.cfg = cfg
self.value = Buffer(torch.ones(1, dtype=torch.float32, device=torch.get_default_device())) self.value = Buffer(torch.ones(1, dtype=torch.float32, device=torch.device('cuda:0')))
self._percentiles = Buffer(torch.tensor([5, 95], dtype=torch.float32, device=torch.get_default_device())) self._percentiles = Buffer(torch.tensor([5, 95], dtype=torch.float32, device=torch.device('cuda:0')))
def state_dict(self): def state_dict(self):
return dict(value=self.value, percentiles=self._percentiles) return dict(value=self.value, percentiles=self._percentiles)

1
tdmpc2/common/world_model.py
View File

@@ -7,6 +7,7 @@ from common import layers, math, init
from tensordict import TensorDict from tensordict import TensorDict
from tensordict.nn import TensorDictParams from tensordict.nn import TensorDictParams
class WorldModel(nn.Module): class WorldModel(nn.Module):
""" """
TD-MPC2 implicit world model architecture. TD-MPC2 implicit world model architecture.

8
tdmpc2/envs/__init__.py
View File

@@ -9,11 +9,8 @@ from envs.wrappers.tensor import TensorWrapper
def missing_dependencies(task): def missing_dependencies(task):
raise ValueError(f'Missing dependencies for task {task}; install dependencies to use this environment.') raise ValueError(f'Missing dependencies for task {task}; install dependencies to use this environment.')
from envs.dmcontrol import make_env as make_dm_control_env
try: try:
pass from envs.dmcontrol import make_env as make_dm_control_env
except: except:
make_dm_control_env = missing_dependencies make_dm_control_env = missing_dependencies
try: try:
@@ -67,8 +64,7 @@ def make_env(cfg):
for fn in [make_dm_control_env, make_maniskill_env, make_metaworld_env, make_myosuite_env]: for fn in [make_dm_control_env, make_maniskill_env, make_metaworld_env, make_myosuite_env]:
try: try:
env = fn(cfg) env = fn(cfg)
except ValueError as err: except ValueError:
print(err)
pass pass
if env is None: if env is None:
raise ValueError(f'Failed to make environment "{cfg.task}": please verify that dependencies are installed and that the task exists.') raise ValueError(f'Failed to make environment "{cfg.task}": please verify that dependencies are installed and that the task exists.')

7
tdmpc2/evaluate.py
View File

@@ -29,7 +29,7 @@ def evaluate(cfg: dict):
`eval_episodes`: number of episodes to evaluate on per task (default: 10) `eval_episodes`: number of episodes to evaluate on per task (default: 10)
`save_video`: whether to save a video of the evaluation (default: True) `save_video`: whether to save a video of the evaluation (default: True)
`seed`: random seed (default: 1) `seed`: random seed (default: 1)
See config.yaml for a full list of args. See config.yaml for a full list of args.
Example usage: Example usage:
@@ -39,8 +39,7 @@ def evaluate(cfg: dict):
$ python evaluate.py task=dog-run checkpoint=/path/to/dog-1.pt save_video=true $ python evaluate.py task=dog-run checkpoint=/path/to/dog-1.pt save_video=true
``` ```
""" """
if torch.get_default_device().type == "cuda": assert torch.cuda.is_available()
assert torch.cuda.is_available()
assert cfg.eval_episodes > 0, 'Must evaluate at least 1 episode.' assert cfg.eval_episodes > 0, 'Must evaluate at least 1 episode.'
cfg = parse_cfg(cfg) cfg = parse_cfg(cfg)
set_seed(cfg.seed) set_seed(cfg.seed)
@@ -58,7 +57,7 @@ def evaluate(cfg: dict):
agent = TDMPC2(cfg) agent = TDMPC2(cfg)
assert os.path.exists(cfg.checkpoint), f'Checkpoint {cfg.checkpoint} not found! Must be a valid filepath.' assert os.path.exists(cfg.checkpoint), f'Checkpoint {cfg.checkpoint} not found! Must be a valid filepath.'
agent.load(cfg.checkpoint) agent.load(cfg.checkpoint)
# Evaluate # Evaluate
if cfg.multitask: if cfg.multitask:
print(colored(f'Evaluating agent on {len(cfg.tasks)} tasks:', 'yellow', attrs=['bold'])) print(colored(f'Evaluating agent on {len(cfg.tasks)} tasks:', 'yellow', attrs=['bold']))

25
tdmpc2/tdmpc2.py
View File

@@ -1,5 +1,3 @@
import os
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
@@ -8,7 +6,6 @@ from common.scale import RunningScale
from common.world_model import WorldModel from common.world_model import WorldModel
from tensordict import TensorDict from tensordict import TensorDict
torch.set_default_device(os.getenv("TDMPC2_DEFAULT_DEVICE", "cuda:0"))
class TDMPC2(torch.nn.Module): class TDMPC2(torch.nn.Module):
""" """
@@ -20,7 +17,7 @@ class TDMPC2(torch.nn.Module):
def __init__(self, cfg): def __init__(self, cfg):
super().__init__() super().__init__()
self.cfg = cfg self.cfg = cfg
self.device = torch.get_default_device() self.device = torch.device('cuda:0')
self.model = WorldModel(cfg).to(self.device) self.model = WorldModel(cfg).to(self.device)
self.optim = torch.optim.Adam([ self.optim = torch.optim.Adam([
{'params': self.model._encoder.parameters(), 'lr': self.cfg.lr*self.cfg.enc_lr_scale}, {'params': self.model._encoder.parameters(), 'lr': self.cfg.lr*self.cfg.enc_lr_scale},
@@ -35,7 +32,7 @@ class TDMPC2(torch.nn.Module):
self.scale = RunningScale(cfg) self.scale = RunningScale(cfg)
self.cfg.iterations += 2*int(cfg.action_dim >= 20) # Heuristic for large action spaces self.cfg.iterations += 2*int(cfg.action_dim >= 20) # Heuristic for large action spaces
self.discount = torch.tensor( self.discount = torch.tensor(
[self._get_discount(ep_len) for ep_len in cfg.episode_lengths], device=torch.get_default_device() [self._get_discount(ep_len) for ep_len in cfg.episode_lengths], device='cuda:0'
) if self.cfg.multitask else self._get_discount(cfg.episode_length) ) if self.cfg.multitask else self._get_discount(cfg.episode_length)
self._prev_mean = torch.nn.Buffer(torch.zeros(self.cfg.horizon, self.cfg.action_dim, device=self.device)) self._prev_mean = torch.nn.Buffer(torch.zeros(self.cfg.horizon, self.cfg.action_dim, device=self.device))
if cfg.compile: if cfg.compile:
@@ -91,16 +88,16 @@ class TDMPC2(torch.nn.Module):
name_map = [ name_map = [
"weight", "bias", "ln.weight", "ln.bias", "weight", "bias", "ln.weight", "ln.bias",
] ]
print("Listing state dict keys (from disk)") # print("Listing state dict keys (from disk)")
for k in list(local_state_dict.keys()): # for k in list(local_state_dict.keys()):
print("\t", k) # print("\t", k)
sd = model.state_dict() sd = model.state_dict()
print("Listing dest state dict keys") # print("Listing dest state dict keys")
for k in list(sd.keys()): # for k in list(sd.keys()):
print("\t", k) # print("\t", k)
print("Maps:") # print("Maps:")
new_sd = dict(sd) new_sd = dict(sd)
for cur_prefix in (prefix, "_target"+prefix[:-1]+"_"): for cur_prefix in (prefix, "_target"+prefix[:-1]+"_"):
for key, val in list(local_state_dict.items()): for key, val in list(local_state_dict.items()):
@@ -109,12 +106,12 @@ class TDMPC2(torch.nn.Module):
num = key[len(cur_prefix + "params."):] num = key[len(cur_prefix + "params."):]
new_key = str(int(num) // 4) + "." + name_map[int(num) % 4] new_key = str(int(num) // 4) + "." + name_map[int(num) % 4]
new_total_key = cur_prefix + 'params.' + new_key new_total_key = cur_prefix + 'params.' + new_key
print("\t", key, '-->', new_total_key) # print("\t", key, '-->', new_total_key)
del local_state_dict[key] del local_state_dict[key]
new_sd[new_total_key] = val new_sd[new_total_key] = val
if not cur_prefix.startswith("_target"): if not cur_prefix.startswith("_target"):
new_total_key = "_detach" + cur_prefix[:-1] + "_" + 'params.' + new_key new_total_key = "_detach" + cur_prefix[:-1] + "_" + 'params.' + new_key
print("\t", 'DETACH', key, '-->', new_total_key) # print("\t", 'DETACH', key, '-->', new_total_key)
new_sd[new_total_key] = val new_sd[new_total_key] = val
local_state_dict.update(new_sd) local_state_dict.update(new_sd)
return local_state_dict return local_state_dict

3
tdmpc2/train.py
View File

@@ -43,8 +43,7 @@ def train(cfg: dict):
$ python train.py task=dog-run steps=7000000 $ python train.py task=dog-run steps=7000000
``` ```
""" """
if torch.get_default_device().type == 'cuda': assert torch.cuda.is_available()
assert torch.cuda.is_available()
assert cfg.steps > 0, 'Must train for at least 1 step.' assert cfg.steps > 0, 'Must train for at least 1 step.'
cfg = parse_cfg(cfg) cfg = parse_cfg(cfg)
set_seed(cfg.seed) set_seed(cfg.seed)