Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+254
@@ -0,0 +1,254 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""A standalone script to run DIAL MPC with Genesis (DexMachina)
|
||||
|
||||
Author: Chaoyi Pan
|
||||
Date: 2025-10-16
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import hydra
|
||||
import loguru
|
||||
import numpy as np
|
||||
import torch
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from spider.config import Config, process_config
|
||||
from spider.interp import get_slice
|
||||
from spider.optimizers.sampling import (
|
||||
make_optimize_fn,
|
||||
make_optimize_once_fn,
|
||||
make_rollout_fn,
|
||||
)
|
||||
from spider.simulators.dexmachina import (
|
||||
copy_sample_state,
|
||||
get_obj_arti_dist,
|
||||
get_obj_pos_dist,
|
||||
get_obj_quat_dist,
|
||||
get_qpos,
|
||||
get_reward,
|
||||
get_terminate,
|
||||
get_terminal_reward,
|
||||
get_trace,
|
||||
load_env_params,
|
||||
load_state,
|
||||
save_env_params,
|
||||
save_state,
|
||||
setup_env,
|
||||
step_env,
|
||||
sync_env,
|
||||
)
|
||||
from spider.viewers import setup_viewer, update_viewer
|
||||
|
||||
|
||||
def main(config: Config):
|
||||
"""Run the SPIDER using Genesis (DexMachina) backend"""
|
||||
# Setup env (ref_data set to None since environment has built-in reference)
|
||||
env = setup_env(config, None)
|
||||
config.nu = env.num_actions
|
||||
# process config, set defaults and derived fields
|
||||
config = process_config(config)
|
||||
if config.max_sim_steps == -1:
|
||||
config.max_sim_steps = env.max_episode_length
|
||||
else:
|
||||
config.max_sim_steps = min(env.max_episode_length, config.max_sim_steps)
|
||||
|
||||
# Create placeholder reference data for compatibility
|
||||
ref_data = (
|
||||
torch.zeros(
|
||||
config.max_sim_steps + config.horizon_steps + config.ctrl_steps,
|
||||
config.nu,
|
||||
device=config.device,
|
||||
),
|
||||
)
|
||||
|
||||
# Setup env params (curriculum for object joint stiffness/damping)
|
||||
env_params_list = []
|
||||
kp_max = 0.0
|
||||
kv_max = 0.0
|
||||
eta = 0.001 ** (1.0 / (config.max_num_iterations - 1))
|
||||
kp_list = kp_max * (eta ** np.arange(config.max_num_iterations))
|
||||
kp_list[-1] = 0.0
|
||||
kv_list = kv_max * (eta ** np.arange(config.max_num_iterations))
|
||||
kv_list[-1] = 0.0
|
||||
kp_env = env.objects[env.object_names[0]].entity.get_dofs_kp()
|
||||
kv_env = env.objects[env.object_names[0]].entity.get_dofs_kv()
|
||||
for i in range(config.max_num_iterations):
|
||||
env_params = []
|
||||
for j in range(config.num_dr):
|
||||
kp = torch.ones_like(kp_env, device=config.device) * kp_list[i]
|
||||
kv = torch.ones_like(kv_env, device=config.device) * kv_list[i]
|
||||
env_params.append({"kp": kp, "kv": kv})
|
||||
env_params_list.append(env_params)
|
||||
config.env_params_list = env_params_list
|
||||
|
||||
# setup viewer
|
||||
run_viewer = setup_viewer(config, None, None)
|
||||
|
||||
# setup optimizer
|
||||
rollout = make_rollout_fn(
|
||||
step_env,
|
||||
save_state,
|
||||
load_state,
|
||||
get_reward,
|
||||
get_terminal_reward,
|
||||
get_terminate,
|
||||
get_trace,
|
||||
save_env_params,
|
||||
load_env_params,
|
||||
copy_sample_state,
|
||||
)
|
||||
optimize_once = make_optimize_once_fn(rollout)
|
||||
optimize = make_optimize_fn(optimize_once)
|
||||
|
||||
# initial controls
|
||||
ctrls = torch.zeros(
|
||||
config.horizon_steps,
|
||||
config.nu,
|
||||
device=config.device,
|
||||
)
|
||||
|
||||
# buffers for saving info and trajectory
|
||||
info_list = []
|
||||
|
||||
# run viewer + control loop
|
||||
t_start = time.perf_counter()
|
||||
sim_step = 0
|
||||
with run_viewer() as viewer:
|
||||
while viewer.is_running():
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# optimize using future reference window at control-rate (+1 lookahead)
|
||||
ref_slice = get_slice(
|
||||
ref_data, sim_step + 1, sim_step + config.horizon_steps + 1
|
||||
)
|
||||
ctrls, infos = optimize(config, env, ctrls, ref_slice)
|
||||
infos["sim_step"] = sim_step
|
||||
|
||||
# step environment for ctrl_steps
|
||||
step_info = {
|
||||
"obj_pos_dist": [],
|
||||
"obj_quat_dist": [],
|
||||
"obj_arti_dist": [],
|
||||
"qpos": [],
|
||||
}
|
||||
for i in range(config.ctrl_steps):
|
||||
ctrl = ctrls[i]
|
||||
ctrl_repeat = ctrl.unsqueeze(0).repeat(
|
||||
int(config.num_samples), 1
|
||||
) # (batch_size, num_actions)
|
||||
step_env(config, env, ctrl_repeat)
|
||||
step_info["obj_pos_dist"].append(get_obj_pos_dist(env)[0].item())
|
||||
step_info["obj_quat_dist"].append(get_obj_quat_dist(env)[0].item())
|
||||
step_info["obj_arti_dist"].append(get_obj_arti_dist(env)[0].item())
|
||||
step_info["qpos"].append(get_qpos(env)[0].detach().cpu().numpy())
|
||||
sim_step += 1
|
||||
if config.save_video:
|
||||
env._recording = True
|
||||
env._render_headless()
|
||||
env._recording = False
|
||||
for k in step_info:
|
||||
step_info[k] = np.stack(step_info[k], axis=0)
|
||||
infos.update(step_info)
|
||||
|
||||
# sync env state
|
||||
env = sync_env(config, env)
|
||||
|
||||
# receding horizon update
|
||||
prev_ctrl = ctrls[config.ctrl_steps :]
|
||||
new_ctrl = torch.zeros(
|
||||
(config.ctrl_steps, config.nu),
|
||||
device=config.device,
|
||||
)
|
||||
ctrls = torch.cat([prev_ctrl, new_ctrl], dim=0)
|
||||
|
||||
# sync viewer state and render
|
||||
# Create dummy mj_data with time for viewer
|
||||
dummy_mj_data = type(
|
||||
"DummyMjData", (), {"time": sim_step * config.sim_dt}
|
||||
)()
|
||||
update_viewer(
|
||||
config,
|
||||
viewer,
|
||||
mj_model=None,
|
||||
mj_data=dummy_mj_data,
|
||||
mj_data_ref=None,
|
||||
info=infos,
|
||||
)
|
||||
|
||||
# progress
|
||||
t1 = time.perf_counter()
|
||||
rtr = config.ctrl_dt / (t1 - t0)
|
||||
print(
|
||||
f"Realtime rate: {rtr:.2f}, plan time: {t1 - t0:.4f}s, sim_steps: {sim_step}/{config.max_sim_steps}, opt_steps: {infos['opt_steps'][0]}",
|
||||
end="\r",
|
||||
)
|
||||
|
||||
# record info/trajectory at control tick
|
||||
# rule out "trace"
|
||||
info_list.append({k: v for k, v in infos.items() if k != "trace_sample"})
|
||||
|
||||
if sim_step >= config.max_sim_steps:
|
||||
break
|
||||
|
||||
t_end = time.perf_counter()
|
||||
print(f"\nTotal time: {t_end - t_start:.4f}s")
|
||||
|
||||
# print object dist in mean and std
|
||||
obj_pos_dists = np.concatenate([info["obj_pos_dist"] for info in info_list])
|
||||
obj_quat_dists = np.concatenate([info["obj_quat_dist"] for info in info_list])
|
||||
obj_arti_dists = np.concatenate([info["obj_arti_dist"] for info in info_list])
|
||||
print(f"obj_pos_dist: {obj_pos_dists.mean():.4f} ± {obj_pos_dists.std():.4f}")
|
||||
print(f"obj_quat_dist: {obj_quat_dists.mean():.4f} ± {obj_quat_dists.std():.4f}")
|
||||
print(f"obj_arti_dist: {obj_arti_dists.mean():.4f} ± {obj_arti_dists.std():.4f}")
|
||||
|
||||
# save retargeted trajectory
|
||||
if config.save_info and len(info_list) > 0:
|
||||
info_aggregated = {}
|
||||
for k in info_list[0].keys():
|
||||
info_aggregated[k] = np.stack([info[k] for info in info_list], axis=0)
|
||||
np.savez(f"{config.output_dir}/trajectory_dexmachina.npz", **info_aggregated)
|
||||
loguru.logger.info(
|
||||
f"Saved info to {config.output_dir}/trajectory_dexmachina.npz"
|
||||
)
|
||||
|
||||
# save video
|
||||
if config.save_video and len(env._recorded_frames) > 0:
|
||||
env.export_video(
|
||||
f"{config.output_dir}/trajectory_dexmachina.mp4", wait_for_max=False
|
||||
)
|
||||
loguru.logger.info(
|
||||
f"Saved video to {config.output_dir}/trajectory_dexmachina.mp4"
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="config", config_name="dexmachina")
|
||||
def run_main(cfg: DictConfig) -> None:
|
||||
# Convert DictConfig to Config dataclass, handling special fields
|
||||
config_dict = dict(cfg)
|
||||
|
||||
# Handle special conversions
|
||||
if "noise_scale" in config_dict and config_dict["noise_scale"] is None:
|
||||
config_dict.pop("noise_scale") # Let the default factory handle it
|
||||
|
||||
# Convert lists to tuples where needed
|
||||
if "pair_margin_range" in config_dict:
|
||||
config_dict["pair_margin_range"] = tuple(config_dict["pair_margin_range"])
|
||||
if "xy_offset_range" in config_dict:
|
||||
config_dict["xy_offset_range"] = tuple(config_dict["xy_offset_range"])
|
||||
|
||||
config = Config(**config_dict)
|
||||
main(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_main()
|
||||
Reference in New Issue
Block a user