383 lines
12 KiB
Python
383 lines
12 KiB
Python
# 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()
|