62 lines
2.7 KiB
Python
62 lines
2.7 KiB
Python
def make_policy(cfg):
|
|
if cfg.policy.name != "diffusion" and cfg.rollout_batch_size > 1:
|
|
raise NotImplementedError("Only diffusion policy supports rollout_batch_size > 1 for the time being.")
|
|
|
|
if cfg.policy.name == "tdmpc":
|
|
from lerobot.common.policies.tdmpc.policy import TDMPC
|
|
|
|
policy = TDMPC(cfg.policy, cfg.device)
|
|
elif cfg.policy.name == "diffusion":
|
|
from lerobot.common.policies.diffusion.policy import DiffusionPolicy
|
|
|
|
policy = DiffusionPolicy(
|
|
cfg=cfg.policy,
|
|
cfg_device=cfg.device,
|
|
cfg_noise_scheduler=cfg.noise_scheduler,
|
|
cfg_rgb_model=cfg.rgb_model,
|
|
cfg_obs_encoder=cfg.obs_encoder,
|
|
cfg_optimizer=cfg.optimizer,
|
|
n_action_steps=cfg.n_action_steps + cfg.n_latency_steps,
|
|
**cfg.policy,
|
|
)
|
|
elif cfg.policy.name == "act":
|
|
from lerobot.common.policies.act.policy import ActionChunkingTransformerPolicy
|
|
|
|
policy = ActionChunkingTransformerPolicy(
|
|
cfg.policy, cfg.device, n_action_steps=cfg.n_action_steps + cfg.n_latency_steps
|
|
)
|
|
else:
|
|
raise ValueError(cfg.policy.name)
|
|
|
|
if cfg.policy.pretrained_model_path:
|
|
# TODO(rcadene): hack for old pretrained models from fowm
|
|
if cfg.policy.name == "tdmpc" and "fowm" in cfg.policy.pretrained_model_path:
|
|
if "offline" in cfg.pretrained_model_path:
|
|
policy.step[0] = 25000
|
|
elif "final" in cfg.pretrained_model_path:
|
|
policy.step[0] = 100000
|
|
else:
|
|
raise NotImplementedError()
|
|
policy.load(cfg.policy.pretrained_model_path)
|
|
|
|
# import torch
|
|
# loaded = torch.load('/home/alexander/Downloads/dp.pth')
|
|
# aligned = {}
|
|
|
|
# their_prefix = "obs_encoder.obs_nets.image.backbone"
|
|
# our_prefix = "obs_encoder.key_model_map.image.backbone"
|
|
# aligned.update({our_prefix + k.removeprefix(their_prefix): v for k, v in loaded.items() if k.startswith(their_prefix)})
|
|
# their_prefix = "obs_encoder.obs_nets.image.pool"
|
|
# our_prefix = "obs_encoder.key_model_map.image.pool"
|
|
# aligned.update({our_prefix + k.removeprefix(their_prefix): v for k, v in loaded.items() if k.startswith(their_prefix)})
|
|
# their_prefix = "obs_encoder.obs_nets.image.nets.3"
|
|
# our_prefix = "obs_encoder.key_model_map.image.out"
|
|
# aligned.update({our_prefix + k.removeprefix(their_prefix): v for k, v in loaded.items() if k.startswith(their_prefix)})
|
|
|
|
# aligned.update({k: v for k, v in loaded.items() if k.startswith('model.')})
|
|
# missing_keys, unexpected_keys = policy.diffusion.load_state_dict(aligned, strict=False)
|
|
# assert all('_dummy_variable' in k for k in missing_keys)
|
|
# assert len(unexpected_keys) == 0
|
|
|
|
return policy
|