9d002032d1
WIP Aloha env tests pass Rendering works (fps look fast tho? TODO action bounding is too wide [-1,1]) Update README Copy past from act repo Remove download.py add a WIP for Simxarm Remove download.py add a WIP for Simxarm Add act yaml (TODO: try train.py) Training can runs (TODO: eval) Add tasks without end_effector that are compatible with dataset, Eval can run (TODO: training and pretrained model) Add AbstractEnv, Refactor AlohaEnv, Add rendering_hook in env, Minor modifications, (TODO: Refactor Pusht and Simxarm) poetry lock fix bug in compute_stats for action normalization fix more bugs in normalization fix training fix import PushtEnv inheriates AbstractEnv, Improve factory Normalization Add _make_env to EnvAbstract Add call_rendering_hooks to pusht env SimxarmEnv inherites from AbstractEnv (NOT TESTED) Add aloha tests artifacts + update pusht stats fix image normalization: before env was in [0,1] but dataset in [0,255], and now both in [0,255] Small fix on simxarm Add next to obs Add top camera to Aloha env (TODO: make it compatible with set of cameras) Add top camera to Aloha env (TODO: make it compatible with set of cameras)
41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
def make_policy(cfg):
|
|
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,
|
|
cfg_ema=cfg.ema,
|
|
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)
|
|
|
|
return policy
|