# 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()