Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+382
@@ -0,0 +1,382 @@
|
||||
# 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.
|
||||
|
||||
"""Run SPIDER MPC with ManipTrans (IsaacGym) backend.
|
||||
|
||||
Usage:
|
||||
# Default (0f900@0, inspire, left hand):
|
||||
python examples/run_maniptrans.py
|
||||
|
||||
# Different data sequence:
|
||||
python examples/run_maniptrans.py task=0f900@0
|
||||
|
||||
# With video recording:
|
||||
python examples/run_maniptrans.py save_video=true
|
||||
|
||||
# Fewer MPC iterations (faster):
|
||||
python examples/run_maniptrans.py max_num_iterations=16
|
||||
|
||||
Author: Chaoyi Pan
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# IsaacGym must be imported before PyTorch (loaded by spider.config)
|
||||
from isaacgym import gymapi as _gymapi # noqa: F401
|
||||
|
||||
import json
|
||||
import os
|
||||
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.maniptrans import (
|
||||
copy_sample_state,
|
||||
get_fingertip_pos_dist,
|
||||
get_joint_pos_dist,
|
||||
get_obj_pos_dist,
|
||||
get_obj_quat_dist,
|
||||
get_qpos,
|
||||
get_reward,
|
||||
get_terminal_reward,
|
||||
get_terminate,
|
||||
get_trace,
|
||||
load_env_params,
|
||||
load_state,
|
||||
record_video,
|
||||
save_env_params,
|
||||
save_state,
|
||||
setup_env,
|
||||
step_env,
|
||||
sync_env,
|
||||
)
|
||||
# Keys in the Hydra config that are ManipTrans-specific and not part of
|
||||
# the Config dataclass. They are extracted before constructing Config and
|
||||
# attached afterwards as extra attributes.
|
||||
_MANIPTRANS_KEYS = [
|
||||
"maniptrans_data_indices",
|
||||
"maniptrans_episode_length",
|
||||
"maniptrans_tighten_method",
|
||||
"maniptrans_tighten_factor",
|
||||
"maniptrans_tighten_steps",
|
||||
"maniptrans_actions_moving_average",
|
||||
"maniptrans_use_pid",
|
||||
"maniptrans_use_quat_rot",
|
||||
"maniptrans_ctrl_freq_inv",
|
||||
"maniptrans_record_video",
|
||||
]
|
||||
|
||||
|
||||
def main(config: Config):
|
||||
"""Run the SPIDER MPC using ManipTrans (IsaacGym) backend."""
|
||||
# Setup env (ref_data set to None since environment has built-in reference)
|
||||
env = setup_env(config, None)
|
||||
# ManipTrans uses base+residual action architecture.
|
||||
# The nominal base actions (DOF targets from demo) are fixed inside step_env.
|
||||
# SPIDER only optimizes the residual (18 dims).
|
||||
config.nu = env.num_actions
|
||||
|
||||
# Save output_dir before process_config (which sets output_dir to
|
||||
# the standard processed directory for MJWP).
|
||||
output_dir = config.output_dir
|
||||
|
||||
# process config: compute steps, noise schedule, output paths
|
||||
config = process_config(config)
|
||||
|
||||
# Restore ManipTrans-specific output_dir if explicitly set
|
||||
if output_dir:
|
||||
config.output_dir = output_dir
|
||||
os.makedirs(config.output_dir, exist_ok=True)
|
||||
|
||||
loguru.logger.info(f"Output directory: {config.output_dir}")
|
||||
loguru.logger.info(
|
||||
f"nu={config.nu}, num_envs={env.num_envs}, "
|
||||
f"horizon_steps={config.horizon_steps}, ctrl_steps={config.ctrl_steps}"
|
||||
)
|
||||
|
||||
# Set max_sim_steps from environment episode length
|
||||
max_episode_length = int(env.demo_data["seq_len"].min().item())
|
||||
if config.max_sim_steps == -1:
|
||||
config.max_sim_steps = max_episode_length
|
||||
else:
|
||||
config.max_sim_steps = min(max_episode_length, config.max_sim_steps)
|
||||
loguru.logger.info(f"max_sim_steps={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 (no curriculum for ManipTrans evaluation)
|
||||
env_params_list = []
|
||||
for i in range(config.max_num_iterations):
|
||||
env_params_list.append([{} for _ in range(config.num_dr)])
|
||||
config.env_params_list = env_params_list
|
||||
|
||||
# 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 = []
|
||||
|
||||
# MPC loop
|
||||
t_start = time.perf_counter()
|
||||
sim_step = 0
|
||||
while sim_step < config.max_sim_steps:
|
||||
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": [],
|
||||
"joint_pos_dist": [],
|
||||
"fingertip_pos_dist": [],
|
||||
"qpos": [],
|
||||
"obj_pos": [],
|
||||
"obj_quat": [],
|
||||
"joint_state": [],
|
||||
"base_state": [],
|
||||
}
|
||||
for i in range(config.ctrl_steps):
|
||||
ctrl = ctrls[i]
|
||||
ctrl_repeat = ctrl.unsqueeze(0).repeat(
|
||||
int(config.num_samples), 1
|
||||
)
|
||||
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["joint_pos_dist"].append(
|
||||
get_joint_pos_dist(env)[0].item()
|
||||
)
|
||||
step_info["fingertip_pos_dist"].append(
|
||||
get_fingertip_pos_dist(env)[0].item()
|
||||
)
|
||||
step_info["qpos"].append(
|
||||
get_qpos(env)[0].detach().cpu().numpy()
|
||||
)
|
||||
step_info["obj_pos"].append(
|
||||
env.states["manip_obj_pos"][0].detach().cpu().numpy()
|
||||
)
|
||||
step_info["obj_quat"].append(
|
||||
env.states["manip_obj_quat"][0].detach().cpu().numpy()
|
||||
)
|
||||
step_info["joint_state"].append(
|
||||
env.states["joints_state"][0, :, :3]
|
||||
.detach()
|
||||
.cpu()
|
||||
.numpy()
|
||||
)
|
||||
step_info["base_state"].append(
|
||||
env.states["base_state"][0].detach().cpu().numpy()
|
||||
)
|
||||
sim_step += 1
|
||||
|
||||
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)
|
||||
|
||||
# progress
|
||||
t1 = time.perf_counter()
|
||||
rtr = config.ctrl_dt / (t1 - t0)
|
||||
print(
|
||||
f" step {sim_step:>4d}/{config.max_sim_steps} | "
|
||||
f"e_t={step_info['obj_pos_dist'].mean():.4f}m | "
|
||||
f"e_r={step_info['obj_quat_dist'].mean():.1f}deg | "
|
||||
f"e_j={step_info['joint_pos_dist'].mean():.4f}m | "
|
||||
f"e_ft={step_info['fingertip_pos_dist'].mean():.4f}m | "
|
||||
f"opt={infos['opt_steps'][0]} | "
|
||||
f"RTR={rtr:.2f}x"
|
||||
)
|
||||
|
||||
# record info/trajectory at control tick
|
||||
info_list.append(
|
||||
{k: v for k, v in infos.items() if k != "trace_sample"}
|
||||
)
|
||||
|
||||
t_end = time.perf_counter()
|
||||
|
||||
# Aggregate all per-step metrics
|
||||
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]
|
||||
)
|
||||
joint_pos_dists = np.concatenate(
|
||||
[info["joint_pos_dist"] for info in info_list]
|
||||
)
|
||||
fingertip_pos_dists = np.concatenate(
|
||||
[info["fingertip_pos_dist"] for info in info_list]
|
||||
)
|
||||
|
||||
# RL-compatible metrics (matching ManipTrans eval_score.py):
|
||||
# e_t = mean object translation error (m)
|
||||
# e_r = mean object rotation error (deg)
|
||||
# e_j = mean joint position error (m)
|
||||
# e_ft = mean fingertip position error (m)
|
||||
e_t = float(obj_pos_dists.mean())
|
||||
e_r = float(obj_quat_dists.mean())
|
||||
e_j = float(joint_pos_dists.mean())
|
||||
e_ft = float(fingertip_pos_dists.mean())
|
||||
|
||||
loguru.logger.info(
|
||||
f"Total time: {t_end - t_start:.1f}s | "
|
||||
f"e_t={e_t * 100:.2f}cm | e_r={e_r:.2f}deg | "
|
||||
f"e_j={e_j * 100:.2f}cm | e_ft={e_ft * 100:.2f}cm"
|
||||
)
|
||||
|
||||
# save evaluation metrics
|
||||
if config.save_info:
|
||||
metrics = {
|
||||
# RL-compatible metrics (same keys as ManipTrans eval.json)
|
||||
"e_t": e_t,
|
||||
"e_r": e_r,
|
||||
"e_j": e_j,
|
||||
"e_ft": e_ft,
|
||||
# Extended statistics
|
||||
"e_t_std": float(obj_pos_dists.std()),
|
||||
"e_r_std": float(obj_quat_dists.std()),
|
||||
"e_j_std": float(joint_pos_dists.std()),
|
||||
"e_ft_std": float(fingertip_pos_dists.std()),
|
||||
# Run info
|
||||
"total_sim_steps": sim_step,
|
||||
"total_time": t_end - t_start,
|
||||
"data_index": config.maniptrans_data_indices[0],
|
||||
"robot_type": config.robot_type,
|
||||
"embodiment_type": config.embodiment_type,
|
||||
}
|
||||
metrics_path = f"{config.output_dir}/metrics_isaac.json"
|
||||
with open(metrics_path, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
loguru.logger.info(f"Saved metrics to {metrics_path}")
|
||||
|
||||
# save retargeted trajectory
|
||||
traj_path = f"{config.output_dir}/rollout_isaac.npz"
|
||||
if config.save_info and len(info_list) > 0:
|
||||
info_aggregated = {}
|
||||
for k in info_list[0].keys():
|
||||
try:
|
||||
info_aggregated[k] = np.stack(
|
||||
[info[k] for info in info_list], axis=0
|
||||
)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
np.savez(traj_path, **info_aggregated)
|
||||
loguru.logger.info(f"Saved trajectory to {traj_path}")
|
||||
|
||||
# Record video by replaying with a separate small environment.
|
||||
# Destroy MPC environment first — IsaacGym GPU pipeline can't have two sims.
|
||||
if config.save_video and config.save_info:
|
||||
video_path = f"{config.output_dir}/rollout_isaac.mp4"
|
||||
record_video(config, video_path, prev_env=env, traj_path=traj_path)
|
||||
loguru.logger.info(f"Saved video to {video_path}")
|
||||
|
||||
return
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="config", config_name="maniptrans")
|
||||
def run_main(cfg: DictConfig) -> None:
|
||||
config_dict = dict(cfg)
|
||||
|
||||
# Extract ManipTrans-specific keys (not in Config dataclass)
|
||||
maniptrans_extras = {}
|
||||
for key in _MANIPTRANS_KEYS:
|
||||
if key in config_dict:
|
||||
maniptrans_extras[key] = config_dict.pop(key)
|
||||
|
||||
# Handle special conversions
|
||||
if "noise_scale" in config_dict and config_dict["noise_scale"] is None:
|
||||
config_dict.pop("noise_scale")
|
||||
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"])
|
||||
|
||||
# Convert list-like OmegaConf objects to plain Python types
|
||||
if "maniptrans_data_indices" in maniptrans_extras:
|
||||
maniptrans_extras["maniptrans_data_indices"] = list(
|
||||
maniptrans_extras["maniptrans_data_indices"]
|
||||
)
|
||||
|
||||
# Use data index as task name (e.g., "0f900@0") instead of generic "maniptrans".
|
||||
# This puts output under .../0f900@0/0/ instead of .../maniptrans/0/.
|
||||
data_indices = maniptrans_extras.get("maniptrans_data_indices", [])
|
||||
if data_indices and config_dict.get("task") == "maniptrans":
|
||||
config_dict["task"] = data_indices[0]
|
||||
|
||||
# If maniptrans_data_indices not explicitly set, derive from task name
|
||||
if "maniptrans_data_indices" not in maniptrans_extras:
|
||||
maniptrans_extras["maniptrans_data_indices"] = [config_dict["task"]]
|
||||
|
||||
config = Config(**config_dict)
|
||||
|
||||
# Attach ManipTrans-specific fields
|
||||
for key, value in maniptrans_extras.items():
|
||||
setattr(config, key, value)
|
||||
config.headless = True
|
||||
|
||||
main(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_main()
|
||||
Reference in New Issue
Block a user