Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+115
@@ -0,0 +1,115 @@
|
||||
# @package _global_
|
||||
defaults:
|
||||
- _self_
|
||||
|
||||
# === TASK CONFIGURATION ===
|
||||
robot_type: xhand # "inspire", "allegro", "g1"
|
||||
embodiment_type: bimanual # "left", "right", "bimanual", "CMU"
|
||||
task: pick_spoon_bowl
|
||||
seed: 0
|
||||
|
||||
# === DATASET CONFIGURATION ===
|
||||
dataset_dir: example_datasets
|
||||
dataset_name: oakink
|
||||
data_id: 0
|
||||
model_path: ""
|
||||
data_path: ""
|
||||
|
||||
# === SIMULATOR CONFIGURATION ===
|
||||
simulator: mjwp # "isaac" | "mujoco" | "mjwp" | "mjwp_cons" | "mjwp_eq" | "mjwp_cons_eq" | "kinematic"
|
||||
device: cuda:0
|
||||
# Simulation timing
|
||||
sim_dt: 0.01 # simulation timestep
|
||||
ctrl_dt: 0.4 # control timestep
|
||||
ref_dt: 0.02 # reference data timestep (1/50.0)
|
||||
render_dt: 0.02 # rendering timestep
|
||||
horizon: 1.6 # planning horizon
|
||||
knot_dt: 0.4 # knot point spacing
|
||||
max_sim_steps: -1 # maximum simulation steps (-1 for unlimited)
|
||||
# Simulation constraints
|
||||
nconmax_per_env: 120 # max contacts per environment
|
||||
njmax_per_env: 350 # max joints per environment
|
||||
# Simulation annealing
|
||||
num_dyn: 1 # number of environments for annealing, used for virtual contact constraint
|
||||
# Domain randomization
|
||||
num_dr: 1 # number of domain randomization groups, used for domain randomization
|
||||
pair_margin_range: [-0.005, 0.005]
|
||||
xy_offset_range: [-0.005, 0.005]
|
||||
perturb_force: 0.0
|
||||
perturb_torque: 0.0
|
||||
contact_guidance: false
|
||||
|
||||
# === OPTIMIZER CONFIGURATION ===
|
||||
# Sampling parameters
|
||||
num_samples: 1024
|
||||
temperature: 0.1
|
||||
max_num_iterations: 32
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 1
|
||||
# Noise scheduling
|
||||
first_ctrl_noise_scale: 0.5
|
||||
last_ctrl_noise_scale: 1.0
|
||||
final_noise_scale: 0.1
|
||||
exploit_ratio: 0.01
|
||||
exploit_noise_scale: 0.01
|
||||
# Noise scaling by component
|
||||
joint_noise_scale: 0.15
|
||||
pos_noise_scale: 0.03
|
||||
rot_noise_scale: 0.03
|
||||
# Reward scaling
|
||||
base_pos_rew_scale: 1.0
|
||||
base_rot_rew_scale: 0.3
|
||||
joint_rew_scale: 0.003
|
||||
pos_rew_scale: 1.0
|
||||
rot_rew_scale: 0.3
|
||||
vel_rew_scale: 0.0001
|
||||
terminal_rew_scale: 1.0
|
||||
# Termination threshold
|
||||
terminate_resample: false
|
||||
object_pos_threshold: 0.1 # meters
|
||||
object_rot_threshold: 0.3 # radians
|
||||
# Gibbs sampling
|
||||
gibbs_sampling: false
|
||||
|
||||
# === VISUALIZATION CONFIGURATION ===
|
||||
show_viewer: true
|
||||
viewer: mujoco-rerun # "mujoco" | "rerun" | "viser" | "isaac"
|
||||
rerun_spawn: true
|
||||
save_video: true
|
||||
save_rerun: false
|
||||
save_viser: false
|
||||
save_metrics: true
|
||||
save_info: true
|
||||
save_config: true
|
||||
|
||||
# === TRACE RECORDING ===
|
||||
trace_dt: 0.02 # 1/50.0
|
||||
num_trace_uniform_samples: 4
|
||||
num_trace_topk_samples: 2
|
||||
trace_site_ids: []
|
||||
|
||||
# === AUTOMATICALLY SET PROPERTIES ===
|
||||
# Computed timesteps
|
||||
horizon_steps: -1
|
||||
knot_steps: -1
|
||||
ref_steps: -1
|
||||
ctrl_steps: -1
|
||||
# Model dimensions
|
||||
nq_obj: -1 # object DOF
|
||||
nq: -1 # total position DOF
|
||||
nv: -1 # total velocity DOF
|
||||
nu: -1 # total control DOF
|
||||
npair: -1 # total pair DOF
|
||||
# Computed tensors (these will be set programmatically)
|
||||
noise_scale: null
|
||||
beta_traj: -1.0
|
||||
# Runtime state
|
||||
env_params_list: []
|
||||
viewer_body_entity_and_ids: []
|
||||
output_dir: ""
|
||||
|
||||
# Hydra configuration
|
||||
hydra:
|
||||
run:
|
||||
dir: .
|
||||
output_subdir: null
|
||||
@@ -0,0 +1,50 @@
|
||||
# Configuration for DexMachina (Genesis) dynamics retargeting pipeline
|
||||
defaults:
|
||||
- _self_
|
||||
|
||||
# Dataset configuration (not used by DexMachina, but kept for compatibility)
|
||||
dataset_dir: "example_datasets"
|
||||
dataset_name: "arctic"
|
||||
robot_type: "inspire_hand" # "inspire_hand", "xhand", "dex3_hand"
|
||||
embodiment_type: "bimanual"
|
||||
task: "box-30-230" # format: "object_name-start_frame-end_frame"
|
||||
data_id: 0
|
||||
viewer: "gs"
|
||||
rerun_spawn: false
|
||||
device: "cuda:0"
|
||||
|
||||
# Simulation and control parameters
|
||||
simulator: dexmachina
|
||||
sim_dt: 0.02 # Genesis simulation timestep
|
||||
ctrl_dt: 0.20 # Control timestep (how often MPC updates)
|
||||
knot_dt: 0.20 # Knot point spacing
|
||||
ref_dt: 0.02 # Reference trajectory timestep
|
||||
horizon: 0.40 # MPC horizon in seconds
|
||||
show_viewer: true
|
||||
save_video: true
|
||||
save_info: true
|
||||
max_sim_steps: -1 # -1 means use full trajectory length
|
||||
nu: 36
|
||||
|
||||
# MPC optimization parameters
|
||||
max_num_iterations: 32 # Number of iterations per MPC step
|
||||
num_samples: 1024 # Number of parallel samples for CEM
|
||||
last_ctrl_noise_scale: 1.0
|
||||
first_ctrl_noise_scale: 0.5
|
||||
final_noise_scale: 0.01
|
||||
pos_noise_scale: 1.0
|
||||
rot_noise_scale: 1.0
|
||||
joint_noise_scale: 1.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 4
|
||||
terminal_rew_scale: 10.0
|
||||
temperature: 1.0
|
||||
|
||||
# Reward mode
|
||||
use_rl_reward: true # true to use dexmachina RL training reward formulation
|
||||
|
||||
# Domain randomization (curriculum learning)
|
||||
num_dr: 1 # Number of domain randomization settings per iteration
|
||||
|
||||
# Output configuration
|
||||
output_dir: "outputs/dexmachina"
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
# Configuration for HDMI (mjlab) humanoid retargeting pipeline
|
||||
defaults:
|
||||
- _self_
|
||||
|
||||
# Task configuration
|
||||
dataset_name: "hdmi"
|
||||
robot_type: "unitree_g1"
|
||||
task: "move_suitcase" # HDMI task name (e.g., move_suitcase, open_drawer, etc.)
|
||||
viewer: "mujoco-rerun" # "mujoco" for visualization, "none" for headless
|
||||
rerun_spawn: false
|
||||
embodiment_type: "humanoid_object" # don't use specific noise scaling for hands
|
||||
device: "cuda:0"
|
||||
joint_noise_scale: 0.2 # disable specific noise scaling for robot
|
||||
|
||||
# Simulation and control parameters
|
||||
simulator: hdmi
|
||||
sim_dt: 0.02 # mjlab simulation timestep
|
||||
ctrl_dt: 0.04 # Control timestep (how often MPC updates)
|
||||
knot_dt: 0.20 # Knot point spacing
|
||||
ref_dt: 0.02 # Reference trajectory timestep
|
||||
horizon: 0.80 # MPC horizon in seconds
|
||||
show_viewer: true
|
||||
save_video: true # Disable for now - need to configure camera names for HDMI
|
||||
save_info: true
|
||||
max_sim_steps: -1 # Maximum simulation steps (-1 for unlimited)
|
||||
nu: -1 # Number of controls (will be set from environment)
|
||||
|
||||
# Optimization parameters
|
||||
use_torch_compile: true
|
||||
|
||||
# MPC optimization parameters
|
||||
max_num_iterations: 32 # Number of iterations per MPC step
|
||||
num_samples: 1024 # Number of parallel samples for sampling
|
||||
last_ctrl_noise_scale: 1.0
|
||||
first_ctrl_noise_scale: 1.0
|
||||
final_noise_scale: 0.1
|
||||
improvement_threshold: 0.02
|
||||
improvement_check_steps: 1
|
||||
terminal_rew_scale: 1.0
|
||||
temperature: 0.1
|
||||
num_resamples: 0
|
||||
resample_ratio: 0.2
|
||||
|
||||
# Domain randomization (not used for HDMI, kept for compatibility)
|
||||
num_dr: 1
|
||||
|
||||
# Output configuration
|
||||
output_dir: "outputs/hdmi"
|
||||
@@ -0,0 +1,62 @@
|
||||
# Configuration for ManipTrans (IsaacGym) dynamics retargeting pipeline
|
||||
defaults:
|
||||
- _self_
|
||||
|
||||
# Dataset configuration
|
||||
dataset_dir: "example_datasets"
|
||||
dataset_name: "maniptrans"
|
||||
robot_type: "inspire" # "inspire", "shadow", "allegro", "xhand"
|
||||
embodiment_type: "left" # "right", "left", "bimanual"
|
||||
task: "0f900@0" # HO-Tracker data index used as task name (format: sequence_id@sub_id)
|
||||
data_id: 0
|
||||
seed: 0
|
||||
viewer: "none" # "isaac" for isaac gym native viewer
|
||||
rerun_spawn: false
|
||||
device: "cuda:0"
|
||||
|
||||
# Simulation and control parameters
|
||||
simulator: maniptrans
|
||||
sim_dt: 0.0166667 # 1/60 second (IsaacGym default)
|
||||
ctrl_dt: 0.2 # Control timestep (how often MPC updates)
|
||||
knot_dt: 0.2 # Knot point spacing
|
||||
ref_dt: 0.0166667 # Reference trajectory timestep
|
||||
render_dt: 0.0166667 # Same as sim_dt
|
||||
trace_dt: 0.0166667 # Must be divisible by sim_dt
|
||||
horizon: 0.6 # MPC horizon in seconds
|
||||
show_viewer: false
|
||||
save_video: true
|
||||
save_info: true
|
||||
max_sim_steps: -1 # -1 means use full trajectory length
|
||||
nu: 18 # 6 (wrist force/torque) + 12 (inspire DOFs), overridden by env
|
||||
|
||||
# MPC optimization parameters
|
||||
max_num_iterations: 32 # Number of iterations per MPC step
|
||||
num_samples: 1024 # Number of parallel samples for CEM
|
||||
last_ctrl_noise_scale: 1.0
|
||||
first_ctrl_noise_scale: 0.5
|
||||
final_noise_scale: 0.1
|
||||
# this part is handled by base controller in maniptrans
|
||||
pos_noise_scale: 0.0
|
||||
rot_noise_scale: 0.0
|
||||
joint_noise_scale: 1.0
|
||||
improvement_threshold: 0.001
|
||||
improvement_check_steps: 4
|
||||
terminal_rew_scale: 10.0
|
||||
temperature: 1.0
|
||||
|
||||
# Domain randomization (not used for ManipTrans evaluation)
|
||||
num_dr: 1
|
||||
|
||||
# ManipTrans-specific parameters (maniptrans_data_indices derived from task name)
|
||||
maniptrans_episode_length: 1200
|
||||
maniptrans_tighten_method: "None" # No curriculum tightening for MPC
|
||||
maniptrans_tighten_factor: 1.0
|
||||
maniptrans_tighten_steps: 3200
|
||||
maniptrans_actions_moving_average: 0.4
|
||||
maniptrans_use_pid: false
|
||||
maniptrans_use_quat_rot: false
|
||||
maniptrans_ctrl_freq_inv: 1
|
||||
maniptrans_record_video: false # Record video after MPC (replay with 2-env camera)
|
||||
|
||||
# Output configuration (empty = use standard path under dataset_dir)
|
||||
output_dir: ""
|
||||
+43
@@ -0,0 +1,43 @@
|
||||
dataset_dir: ../../example_datasets
|
||||
dataset_name: oakink
|
||||
robot_type: allegro
|
||||
embodiment_type: bimanual
|
||||
task: pick_spoon_bowl
|
||||
data_id: 0
|
||||
sim_dt: 0.01
|
||||
ctrl_dt: 0.2
|
||||
ref_dt: 0.02
|
||||
show_viewer: true
|
||||
save_video: true
|
||||
save_info: true
|
||||
# Configuration for MJWP (MuJoCo + Warp) dynamics retargeting pipeline
|
||||
defaults:
|
||||
- _self_
|
||||
- launcher: basic
|
||||
|
||||
# Configuration for MJWP (MuJoCo + Warp) dynamics retargeting pipeline
|
||||
dataset_dir: "../example_datasets"
|
||||
dataset_name: "oakink"
|
||||
robot_type: "allegro"
|
||||
embodiment_type: "bimanual"
|
||||
all_tasks: null # if null, process all available tasks, otherwise specify list of tasks
|
||||
overwrite: false # whether to overwrite existing mjwp trajectories
|
||||
viewer: ""
|
||||
device: "cuda:0"
|
||||
method: "mjwp"
|
||||
|
||||
# MJWP-specific optimization parameters (optional overrides)
|
||||
# Set to null to use defaults from Config class
|
||||
horizon: null # default: 1.2
|
||||
sim_dt: null # default: 0.01
|
||||
ctrl_dt: null # default: 0.2
|
||||
show_viewer: true # typically false for batch processing
|
||||
save_video: true
|
||||
save_info: true
|
||||
max_sim_steps: -1 # -1 means use full trajectory length
|
||||
|
||||
# Hydra configuration
|
||||
hydra:
|
||||
run:
|
||||
dir: .
|
||||
output_subdir: null
|
||||
@@ -0,0 +1,22 @@
|
||||
# @package _global_
|
||||
dataset_name: arcticv2
|
||||
task: s01-box_use_01
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 16
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
@@ -0,0 +1,36 @@
|
||||
# @package _global_
|
||||
dataset_name: arcticv2
|
||||
task: s01-box_use_01
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: -1.0
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 4
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
object_rot_threshold: 2.5
|
||||
max_local_retries: 3
|
||||
# Arctic clips start ~2 s before contact (hand approaching, object resting on
|
||||
# the table only via friction). The default sanity check holds ctrl_ref[0]
|
||||
# for 3 s and requires the object not to drift — that's incompatible with a
|
||||
# pre-grasp start frame, so disable it here.
|
||||
sanity_check_seconds: 0.0
|
||||
# Arctic objects (box / laptop / mixer / notebook / waffleiron) decompose
|
||||
# into 25-30 convex parts each. With bimanual xhand the contact + joint-row
|
||||
# counts overflow the default per-env limits ("nefc overflow - please
|
||||
# increase njmax to ~420"). Bump both above the worst observed value with
|
||||
# headroom.
|
||||
nconmax_per_env: 256
|
||||
njmax_per_env: 768
|
||||
@@ -0,0 +1,25 @@
|
||||
# @package _global_
|
||||
dataset_name: fair_fre
|
||||
task: pick_up
|
||||
robot_type: xhand
|
||||
embodiment_type: right
|
||||
data_id: 4
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
ref_dt: 0.02
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 3
|
||||
terminal_rew_scale: 10.0
|
||||
pos_rew_scale: 1.0
|
||||
max_num_iterations: 16
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
@@ -0,0 +1,24 @@
|
||||
# @package _global_
|
||||
# Override configuration for fair_mon dataset
|
||||
dataset_name: fair_mon
|
||||
robot_type: allegro
|
||||
viewer: rerun-mujoco
|
||||
embodiment_type: right
|
||||
task: cat
|
||||
rerun_spawn: true
|
||||
ref_dt: 0.05
|
||||
max_sim_steps: -1
|
||||
horizon: 0.8
|
||||
ctrl_dt: 0.1
|
||||
knot_dt: 0.2
|
||||
num_dr: 1
|
||||
perturb_torque: 0.0
|
||||
pos_noise_scale: 0.05
|
||||
rot_noise_scale: 0.05
|
||||
joint_noise_scale: 0.1
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 16
|
||||
temperature: 1.0
|
||||
# original reference is bad, so we put smaller scale on tracking wrist pose
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
@@ -0,0 +1,21 @@
|
||||
# @package _global_
|
||||
# Override configuration for gigahand dataset
|
||||
dataset_name: gigahand
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
task: p36-tea
|
||||
temperature: 1.0
|
||||
data_id: 0
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 0.5
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 1.0
|
||||
last_ctrl_noise_scale: 2.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 2
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.03
|
||||
pos_noise_scale: 0.03
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
@@ -0,0 +1,33 @@
|
||||
# @package _global_
|
||||
task: p36-tea
|
||||
dataset_name: gigahand
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
first_ctrl_noise_scale: 1.0
|
||||
final_noise_scale: 0.3
|
||||
sim_dt: 0.005
|
||||
max_num_iterations: 32
|
||||
gibbs_sampling: false
|
||||
# guidance related
|
||||
improvement_threshold: 0.0
|
||||
contact_guidance: true
|
||||
guidance_decay_ratio: 0.5
|
||||
init_pos_actuator_gain: 2.0
|
||||
init_pos_actuator_bias: 2.0
|
||||
init_rot_actuator_gain: 0.3
|
||||
init_rot_actuator_bias: 0.3
|
||||
# with guidance, exploration noise can be higher
|
||||
joint_noise_scale: 0.15
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
# with guidance, horizon can be shorter
|
||||
horizon: 1.6
|
||||
knot_dt: 0.4
|
||||
ctrl_dt: 0.2
|
||||
base_pos_rew_scale: 0.03
|
||||
base_rot_rew_scale: 0.03
|
||||
pos_rew_scale: 1.0
|
||||
rot_rew_scale: 1.0
|
||||
@@ -0,0 +1,23 @@
|
||||
# @package _global_
|
||||
dataset_name: gigahand
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
task: p36-tea
|
||||
temperature: 1.0
|
||||
data_id: 0
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 0.5
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 1.0
|
||||
last_ctrl_noise_scale: 2.0
|
||||
improvement_threshold: -1.0
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 8
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.03
|
||||
pos_noise_scale: 0.03
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
object_rot_threshold: 0.5
|
||||
max_local_retries: 3
|
||||
@@ -0,0 +1,16 @@
|
||||
# @package _global_
|
||||
# Override configuration for gigahand dataset
|
||||
dataset_name: gigahand
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
task: p36-tea
|
||||
data_id: 10
|
||||
horizon: 1.6
|
||||
temperature: 1.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 1
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.03
|
||||
pos_noise_scale: 0.03
|
||||
rot_noise_scale: 0.01
|
||||
@@ -0,0 +1,31 @@
|
||||
# @package _global_
|
||||
# Override configuration for fair_mon dataset
|
||||
dataset_name: amass
|
||||
robot_type: unitree_g1
|
||||
embodiment_type: humanoid
|
||||
task: sprint
|
||||
rerun_spawn: true
|
||||
sim_dt: 0.0166667
|
||||
ref_dt: 0.0333333
|
||||
trace_dt: 0.0333333
|
||||
max_sim_steps: -1
|
||||
horizon: 0.8
|
||||
ctrl_dt: 0.1
|
||||
knot_dt: 0.1
|
||||
# annealing parameters
|
||||
first_ctrl_noise_scale: 0.5
|
||||
last_ctrl_noise_scale: 1.0
|
||||
final_noise_scale: 0.1
|
||||
joint_noise_scale: 0.1
|
||||
# sampling
|
||||
improvement_threshold: 0.01
|
||||
max_num_iterations: 16
|
||||
temperature: 0.6
|
||||
num_samples: 2048
|
||||
# reward
|
||||
vel_rew_scale: 0.0
|
||||
joint_rew_scale: 1.0
|
||||
pos_rew_scale: 3.0
|
||||
rot_rew_scale: 1.0
|
||||
contact_rew_scale: 1.0
|
||||
# note: above parameter is optimized for speed. you can have better performance with horizon=1.0
|
||||
@@ -0,0 +1,33 @@
|
||||
# @package _global_
|
||||
# Override configuration for fair_mon dataset
|
||||
dataset_name: omomo
|
||||
robot_type: unitree_g1
|
||||
viewer: mujoco-rerun
|
||||
embodiment_type: humanoid_object
|
||||
task: move_largebox
|
||||
rerun_spawn: true
|
||||
sim_dt: 0.0166667
|
||||
ref_dt: 0.0333333
|
||||
trace_dt: 0.0333333
|
||||
max_sim_steps: -1
|
||||
horizon: 0.8
|
||||
ctrl_dt: 0.1
|
||||
knot_dt: 0.1
|
||||
# annealing parameters
|
||||
first_ctrl_noise_scale: 0.5
|
||||
last_ctrl_noise_scale: 1.0
|
||||
final_noise_scale: 0.1
|
||||
joint_noise_scale: 0.1
|
||||
# sampling
|
||||
improvement_threshold: 0.01
|
||||
max_num_iterations: 16
|
||||
temperature: 1.0
|
||||
num_samples: 2048
|
||||
# reward
|
||||
vel_rew_scale: 0.0
|
||||
joint_rew_scale: 1.0
|
||||
base_pos_rew_scale: 1.0
|
||||
base_rot_rew_scale: 1.0
|
||||
pos_rew_scale: 1.0
|
||||
rot_rew_scale: 0.3
|
||||
# note: above parameter is optimized for speed. you can have better performance with horizon=1.0
|
||||
@@ -0,0 +1,21 @@
|
||||
# @package _global_
|
||||
task: lift_board
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 16
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
@@ -0,0 +1,32 @@
|
||||
# @package _global_
|
||||
task: lift_board
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
first_ctrl_noise_scale: 1.0
|
||||
final_noise_scale: 0.3
|
||||
sim_dt: 0.005
|
||||
max_num_iterations: 32
|
||||
gibbs_sampling: false
|
||||
# guidance related
|
||||
improvement_threshold: 0.0
|
||||
contact_guidance: true
|
||||
guidance_decay_ratio: 0.8
|
||||
init_pos_actuator_gain: 10.0
|
||||
init_pos_actuator_bias: 10.0
|
||||
init_rot_actuator_gain: 0.3
|
||||
init_rot_actuator_bias: 0.3
|
||||
# with guidance, exploration noise can be higher
|
||||
joint_noise_scale: 0.02
|
||||
pos_noise_scale: 0.005
|
||||
rot_noise_scale: 0.01
|
||||
# with guidance, horizon can be shorter
|
||||
horizon: 0.4
|
||||
knot_dt: 0.1
|
||||
ctrl_dt: 0.05
|
||||
base_pos_rew_scale: 0.3
|
||||
base_rot_rew_scale: 0.3
|
||||
pos_rew_scale: 1.0
|
||||
rot_rew_scale: 1.0
|
||||
@@ -0,0 +1,23 @@
|
||||
# @package _global_
|
||||
task: lift_board
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: -1.0
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 8
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
object_rot_threshold: 0.5
|
||||
max_local_retries: 3
|
||||
@@ -0,0 +1,17 @@
|
||||
# @package _global_
|
||||
task: lift_board
|
||||
robot_type: xhand
|
||||
embodiment_type: bimanual
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
horizon: 2.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 32
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
# rot_rew_scale: 0.3
|
||||
joint_noise_scale: 0.03
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
@@ -0,0 +1,22 @@
|
||||
# @package _global_
|
||||
dataset_name: oakinkv2
|
||||
task: pick_spoon_bowl
|
||||
robot_type: xhand
|
||||
embodiment_type: right
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: 0.01
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 16
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
@@ -0,0 +1,24 @@
|
||||
# @package _global_
|
||||
dataset_name: oakinkv2
|
||||
task: pick_spoon_bowl
|
||||
robot_type: xhand
|
||||
embodiment_type: right
|
||||
data_id: 0
|
||||
temperature: 1.0
|
||||
num_samples: 1024
|
||||
ctrl_dt: 1.5
|
||||
knot_dt: 1.0
|
||||
horizon: 3.0
|
||||
first_ctrl_noise_scale: 2.0
|
||||
last_ctrl_noise_scale: 4.0
|
||||
improvement_threshold: -1.0
|
||||
improvement_check_steps: 2
|
||||
max_num_iterations: 4
|
||||
base_pos_rew_scale: 0.1
|
||||
base_rot_rew_scale: 0.03
|
||||
joint_noise_scale: 0.1
|
||||
pos_noise_scale: 0.01
|
||||
rot_noise_scale: 0.01
|
||||
final_noise_scale: 0.01
|
||||
object_rot_threshold: 1.0
|
||||
max_local_retries: 3
|
||||
+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()
|
||||
+328
@@ -0,0 +1,328 @@
|
||||
# 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.
|
||||
|
||||
"""Retargeting for humand-object interaction with HDMI simulator.
|
||||
|
||||
Author: Chaoyi Pan
|
||||
Date: 2025-10-18
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import hydra
|
||||
import imageio
|
||||
import loguru
|
||||
import mujoco
|
||||
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.hdmi import (
|
||||
copy_sample_state,
|
||||
get_reference,
|
||||
get_reward,
|
||||
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 (
|
||||
log_frame,
|
||||
render_image,
|
||||
setup_renderer,
|
||||
setup_viewer,
|
||||
update_viewer,
|
||||
)
|
||||
|
||||
|
||||
def main(config: Config):
|
||||
"""Run the SPIDER using HDMI backend."""
|
||||
# Setup env (ref_data set to None since environment has built-in reference)
|
||||
env = setup_env(config, None)
|
||||
if config.max_sim_steps == -1:
|
||||
config.max_sim_steps = env.max_episode_length
|
||||
loguru.logger.info(f"Max simulation steps set to {config.max_sim_steps}")
|
||||
|
||||
config.nu = env.action_spec.shape[-1]
|
||||
|
||||
# Process config, set defaults and derived fields
|
||||
config = process_config(config)
|
||||
|
||||
# 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 (empty for HDMI, no domain randomization)
|
||||
env_params_list = []
|
||||
for _ in range(config.max_num_iterations):
|
||||
env_params = [{}] * config.num_dr
|
||||
env_params_list.append(env_params)
|
||||
config.env_params_list = env_params_list
|
||||
|
||||
# Get reference data (states and controls)
|
||||
qpos_ref, qvel_ref, ctrl_ref = get_reference(config, env)
|
||||
np.savez(
|
||||
f"{config.output_dir}/trajectory_kinematic.npz",
|
||||
qpos=qpos_ref.detach().cpu().numpy(),
|
||||
qvel=qvel_ref.detach().cpu().numpy(),
|
||||
ctrl=ctrl_ref.detach().cpu().numpy(),
|
||||
)
|
||||
# optional: also save env xml
|
||||
# env.scene.to_zip(Path(config.output_dir) / "../scene.zip")
|
||||
|
||||
# Setup mujoco model and data from HDMI env (for rendering)
|
||||
default_xml_path = Path(
|
||||
"example_datasets/processed/hdmi/unitree_g1/humanoid_object/move_suitcase/scene/mjlab scene.xml"
|
||||
)
|
||||
if default_xml_path.exists():
|
||||
mj_model = mujoco.MjModel.from_xml_path(str(default_xml_path))
|
||||
config.model_path = default_xml_path
|
||||
else:
|
||||
loguru.logger.warning(
|
||||
f"Default XML path {default_xml_path} does not exist, using env model"
|
||||
)
|
||||
mj_model = env.sim.mj_model
|
||||
config.model_path = "hdmi_scene_from_env"
|
||||
mj_data = mujoco.MjData(mj_model)
|
||||
mj_data_ref = mujoco.MjData(mj_model)
|
||||
|
||||
# Initialize mj_data with current env state
|
||||
sim_data = env.sim.data
|
||||
mj_data.qpos[:] = sim_data.qpos[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = sim_data.qvel[0].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data)
|
||||
mj_data.time = 0.0
|
||||
|
||||
# Initialize reference mj_data
|
||||
mj_data_ref.qpos[:] = qpos_ref[0].detach().cpu().numpy()
|
||||
mj_data_ref.qvel[:] = qvel_ref[0].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data_ref)
|
||||
|
||||
# Setup for video rendering
|
||||
images = []
|
||||
|
||||
# Setup viewer and renderer
|
||||
# Note: mjlab viewer is disabled via cfg.viewer.headless in setup_env
|
||||
# For HDMI with rerun, build scene directly from spec and model
|
||||
# if "rerun" in config.viewer:
|
||||
# # Build and log 3D scene from HDMI's spec and model
|
||||
# if default_xml_path.exists():
|
||||
# xml_path = default_xml_path
|
||||
# else:
|
||||
# loguru.logger.info(
|
||||
# f"Default scene path {default_xml_path} does not exist, using None"
|
||||
# )
|
||||
# xml_path = None
|
||||
# loguru.logger.info("Building Rerun scene from HDMI spec and model...")
|
||||
# config.viewer_body_entity_and_ids = build_and_log_scene_from_spec(
|
||||
# spec=mj_spec,
|
||||
# model=mj_model,
|
||||
# xml_path=xml_path,
|
||||
# entity_root="mujoco",
|
||||
# )
|
||||
# loguru.logger.info(
|
||||
# f"Rerun scene built with {len(config.viewer_body_entity_and_ids)} body entities"
|
||||
# )
|
||||
# # Set model_path to dummy value to indicate scene is already built
|
||||
# config.model_path = "hdmi_scene_from_spec"
|
||||
|
||||
run_viewer = setup_viewer(config, mj_model, mj_data)
|
||||
renderer = setup_renderer(config, mj_model)
|
||||
|
||||
# Setup optimizer
|
||||
rollout = make_rollout_fn(
|
||||
step_env,
|
||||
save_state,
|
||||
load_state,
|
||||
get_reward,
|
||||
get_terminal_reward,
|
||||
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 - first horizon_steps from reference
|
||||
ctrls = ctrl_ref[: config.horizon_steps]
|
||||
|
||||
# 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
|
||||
)
|
||||
if config.max_num_iterations > 0:
|
||||
ctrls, infos = optimize(config, env, ctrls, ref_slice)
|
||||
else:
|
||||
infos = {"opt_steps": [0], "improvement": 0.0}
|
||||
infos["sim_step"] = sim_step
|
||||
|
||||
# Step environment for ctrl_steps
|
||||
step_info = {"qpos": [], "qvel": [], "time": [], "ctrl": [], "ctrl_ref": []}
|
||||
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)
|
||||
|
||||
# Update mj_data with current state
|
||||
sim_data = env.sim.data
|
||||
mj_data.qpos[:] = sim_data.qpos[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = sim_data.qvel[0].detach().cpu().numpy()
|
||||
mj_data.time = (sim_step + 1) * config.sim_dt
|
||||
|
||||
# Render video if enabled
|
||||
should_render = (
|
||||
config.save_video
|
||||
and renderer is not None
|
||||
and i % int(np.round(config.render_dt / config.sim_dt)) == 0
|
||||
)
|
||||
if should_render:
|
||||
# Get reference state from reference data
|
||||
ref_idx = min(sim_step, qpos_ref.shape[0] - 1)
|
||||
mj_data_ref.qpos[:] = qpos_ref[ref_idx].detach().cpu().numpy()
|
||||
mj_data_ref.qvel[:] = qvel_ref[ref_idx].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data_ref)
|
||||
image = render_image(
|
||||
config, renderer, mj_model, mj_data, mj_data_ref
|
||||
)
|
||||
images.append(image)
|
||||
if "rerun" in config.viewer or "viser" in config.viewer:
|
||||
# manually log the state
|
||||
log_frame(
|
||||
mj_data,
|
||||
sim_time=mj_data.time,
|
||||
viewer_body_entity_and_ids=config.viewer_body_entity_and_ids,
|
||||
)
|
||||
|
||||
# Record state info
|
||||
step_info["qpos"].append(mj_data.qpos.copy())
|
||||
step_info["qvel"].append(mj_data.qvel.copy())
|
||||
step_info["time"].append(mj_data.time)
|
||||
step_info["ctrl"].append(ctrl.detach().cpu().numpy())
|
||||
step_info["ctrl_ref"].append(ctrl_ref[sim_step].detach().cpu().numpy())
|
||||
|
||||
sim_step += 1
|
||||
|
||||
# Stack step info
|
||||
for k in step_info:
|
||||
step_info[k] = np.stack(step_info[k], axis=0)
|
||||
infos.update(step_info)
|
||||
|
||||
# Sync env state (broadcast from first env to all)
|
||||
env = sync_env(config, env)
|
||||
|
||||
# Receding horizon update
|
||||
prev_ctrl = ctrls[config.ctrl_steps :]
|
||||
new_ctrl = ctrl_ref[
|
||||
sim_step + prev_ctrl.shape[0] : sim_step
|
||||
+ prev_ctrl.shape[0]
|
||||
+ config.ctrl_steps
|
||||
]
|
||||
ctrls = torch.cat([prev_ctrl, new_ctrl], dim=0)
|
||||
|
||||
# Sync viewer state and render
|
||||
sim_data = env.sim.data
|
||||
mj_data.qpos[:] = sim_data.qpos[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = sim_data.qvel[0].detach().cpu().numpy()
|
||||
# Update reference state
|
||||
ref_idx = min(sim_step, qpos_ref.shape[0] - 1)
|
||||
mj_data_ref.qpos[:] = qpos_ref[ref_idx].detach().cpu().numpy()
|
||||
mj_data_ref.qvel[:] = qvel_ref[ref_idx].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data_ref)
|
||||
update_viewer(config, viewer, mj_model, mj_data, mj_data_ref, 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
|
||||
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")
|
||||
|
||||
# Save retargeted trajectory
|
||||
if config.save_info and len(info_list) > 0:
|
||||
info_aggregated = {}
|
||||
for k in info_list[0]:
|
||||
info_aggregated[k] = np.stack([info[k] for info in info_list], axis=0)
|
||||
np.savez(f"{config.output_dir}/trajectory_hdmi.npz", **info_aggregated)
|
||||
loguru.logger.info(f"Saved info to {config.output_dir}/trajectory_hdmi.npz")
|
||||
|
||||
# Save video
|
||||
if config.save_video and len(images) > 0:
|
||||
video_path = f"{config.output_dir}/visualization_hdmi.mp4"
|
||||
imageio.mimsave(
|
||||
video_path,
|
||||
images,
|
||||
fps=int(1 / config.render_dt),
|
||||
)
|
||||
loguru.logger.info(f"Saved video to {video_path}")
|
||||
|
||||
return
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="config", config_name="hdmi")
|
||||
def run_main(cfg: DictConfig) -> None:
|
||||
"""Main entry point for HDMI retargeting."""
|
||||
# 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()
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
#! /bin/zsh
|
||||
export DISPLAY=:1
|
||||
HDMI_DIR="/home/pcy/Research/code/HDMI"
|
||||
|
||||
# source HDMI env
|
||||
source ${HDMI_DIR}/.venv/bin/activate
|
||||
|
||||
# Install required packages
|
||||
# uv pip install opencv-python rerun-sdk loguru hydra-core omegaconf
|
||||
|
||||
# run SPIDER
|
||||
# viewer options:
|
||||
# - "mjlab" for mjlab viewer only
|
||||
# - "mujoco" for mujoco viewer only
|
||||
# - "rerun" for rerun viewer only
|
||||
# - "mujoco-rerun" for both mujoco and rerun (recommended for SPIDER)
|
||||
# - "mjlab-mujoco-rerun" for all three viewers (mjlab may conflict)
|
||||
# python examples/run_hdmi.py task=move_largebox viewer="mujoco-rerun" rerun_spawn=true horizon=0.8 knot_dt=0.2 ctrl_dt=0.04 num_samples=1024 temperature=0.1 max_num_iterations=32 joint_noise_scale=0.2 first_ctrl_noise_scale=1.0 improvement_check_steps=1 improvement_threshold=0.02 max_sim_steps=100 #ctrl_dt=0.02 knot_dt=0.02 horizon=0.02 joint_noise_scale=0.0001 max_num_iterations=1
|
||||
|
||||
python examples/run_hdmi.py task=move_suitcase joint_noise_scale=0.2 knot_dt=0.2 ctrl_dt=0.04 horizon=0.8 +data_id=1 viewer="mujoco-rerun" rerun_spawn=true +save_rerun=true +save_metrics=false max_sim_steps=-1
|
||||
+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()
|
||||
+640
@@ -0,0 +1,640 @@
|
||||
# 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 Mujoco + Warp
|
||||
|
||||
Up to now, domain randomization is not supported. Will add it later.
|
||||
|
||||
Author: Chaoyi Pan
|
||||
Date: 2025-08-11
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import fields
|
||||
from pathlib import Path
|
||||
|
||||
import hydra
|
||||
import imageio
|
||||
import loguru
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
import warp as wp
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from spider.config import (
|
||||
Config,
|
||||
filter_config_fields,
|
||||
load_config_yaml,
|
||||
process_config,
|
||||
)
|
||||
from spider.interp import get_slice
|
||||
from spider.io import load_data
|
||||
from spider.optimizers.sampling import (
|
||||
make_optimize_fn,
|
||||
make_optimize_once_fn,
|
||||
make_rollout_fn,
|
||||
)
|
||||
from spider.postprocess.get_success_rate import compute_object_tracking_error
|
||||
from spider.simulators.mjwp import (
|
||||
_initial_state_sanity_check,
|
||||
compute_contact_point_delta,
|
||||
copy_sample_state,
|
||||
get_qpos,
|
||||
get_qvel,
|
||||
get_reward,
|
||||
get_terminal_reward,
|
||||
get_terminate,
|
||||
get_trace,
|
||||
load_env_params,
|
||||
load_state,
|
||||
save_env_params,
|
||||
save_state,
|
||||
setup_env,
|
||||
setup_mj_model, # mjwp specific
|
||||
step_env,
|
||||
sync_env,
|
||||
)
|
||||
from spider.viewers import (
|
||||
log_frame,
|
||||
render_image,
|
||||
setup_renderer,
|
||||
setup_viewer,
|
||||
update_viewer,
|
||||
)
|
||||
|
||||
_CONFIG_SKIP_FIELDS = {
|
||||
"noise_scale",
|
||||
"env_params_list",
|
||||
"viewer_body_entity_and_ids",
|
||||
}
|
||||
|
||||
|
||||
def _parse_override_tokens(tokens: list[str]) -> dict:
|
||||
allowed = {field.name for field in fields(Config)}
|
||||
override_dict: dict = {}
|
||||
for item in tokens:
|
||||
if "=" not in item:
|
||||
continue
|
||||
key, value = item.split("=", 1)
|
||||
key = key.lstrip("+")
|
||||
if key not in allowed:
|
||||
continue
|
||||
parsed = OmegaConf.to_container(
|
||||
OmegaConf.from_dotlist([f"{key}={value}"]), resolve=True
|
||||
)
|
||||
if isinstance(parsed, dict) and key in parsed:
|
||||
override_dict[key] = parsed[key]
|
||||
return override_dict
|
||||
|
||||
|
||||
def _extract_cli_overrides(cfg: DictConfig) -> dict:
|
||||
"""Extract CLI overrides so they can be applied on top of a loaded config."""
|
||||
overrides = OmegaConf.select(cfg, "hydra.overrides.task") or []
|
||||
override_dict = _parse_override_tokens(overrides)
|
||||
if override_dict:
|
||||
return override_dict
|
||||
return _parse_override_tokens(sys.argv[1:])
|
||||
|
||||
|
||||
def _assert_object_actuator_gains_zero(
|
||||
env, config: Config, stage: str, atol: float = 1e-4
|
||||
) -> None:
|
||||
if not config.contact_guidance or not config.object_actuator_ids:
|
||||
return
|
||||
actuator_ids = np.asarray(config.object_actuator_ids, dtype=int)
|
||||
if not hasattr(env, "model_wp") or not hasattr(env.model_wp, "actuator_gainprm"):
|
||||
raise AssertionError("MJWarp model does not expose actuator_gainprm.")
|
||||
gainprm = wp.to_torch(env.model_wp.actuator_gainprm).detach().cpu().numpy()
|
||||
biasprm = wp.to_torch(env.model_wp.actuator_biasprm).detach().cpu().numpy()
|
||||
if gainprm.ndim == 3:
|
||||
gainprm = gainprm[0]
|
||||
if biasprm.ndim == 3:
|
||||
biasprm = biasprm[0]
|
||||
kp = gainprm[actuator_ids, 0]
|
||||
kd = -biasprm[actuator_ids, 1]
|
||||
assert np.allclose(kp, 0.0, atol=atol), (
|
||||
f"Object actuator Kp not near zero at {stage}: max={np.max(np.abs(kp))}"
|
||||
)
|
||||
assert np.allclose(kd, 0.0, atol=atol), (
|
||||
f"Object actuator Kd not near zero at {stage}: max={np.max(np.abs(kd))}"
|
||||
)
|
||||
|
||||
|
||||
def _normalize_yaml_value(value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.detach().cpu().tolist()
|
||||
if isinstance(value, np.ndarray):
|
||||
return value.tolist()
|
||||
if isinstance(value, np.generic):
|
||||
return value.item()
|
||||
if isinstance(value, tuple):
|
||||
return list(value)
|
||||
return value
|
||||
|
||||
|
||||
def _save_config_yaml(config: Config) -> None:
|
||||
if not config.save_config:
|
||||
return
|
||||
config_dict = {}
|
||||
for field in fields(config):
|
||||
if field.name in _CONFIG_SKIP_FIELDS:
|
||||
continue
|
||||
config_dict[field.name] = _normalize_yaml_value(getattr(config, field.name))
|
||||
output_path = (
|
||||
Path(config.output_dir)
|
||||
/ f"config{'_act' if config.contact_guidance else ''}.yaml"
|
||||
)
|
||||
OmegaConf.save(config=OmegaConf.create(config_dict), f=str(output_path))
|
||||
loguru.logger.info(f"Saved config to {output_path}")
|
||||
|
||||
|
||||
def _get_bimanual_hand_indices(config: Config) -> tuple[list[int], list[int]]:
|
||||
robot_nu = int(config.nu)
|
||||
if config.contact_guidance:
|
||||
obj_dims = (
|
||||
int(config.object_action_dims) if config.object_action_dims > 0 else 12
|
||||
)
|
||||
robot_nu = max(robot_nu - obj_dims, 0)
|
||||
half = robot_nu // 2
|
||||
right_ids = list(range(0, half))
|
||||
left_ids = list(range(half, robot_nu))
|
||||
return right_ids, left_ids
|
||||
|
||||
|
||||
def _apply_noise_mask(
|
||||
base_noise_scale: torch.Tensor, zero_indices: list[int]
|
||||
) -> torch.Tensor:
|
||||
noise_scale = base_noise_scale.clone()
|
||||
if zero_indices:
|
||||
idx = torch.as_tensor(
|
||||
zero_indices, device=base_noise_scale.device, dtype=torch.long
|
||||
)
|
||||
noise_scale[:, :, idx] *= 0.0
|
||||
return noise_scale
|
||||
|
||||
|
||||
def main(config: Config):
|
||||
"""Run the SPIDER using MuJoCo Warp backend"""
|
||||
# process config, set defaults and derived fields
|
||||
config = process_config(config)
|
||||
if config.contact_guidance and config.improvement_threshold > 0.0:
|
||||
loguru.logger.warning(
|
||||
"contact_guidance requires improvement_threshold <= 0; overriding to 0.0."
|
||||
)
|
||||
config.improvement_threshold = 0.0
|
||||
|
||||
# load reference data (already interpolated and extended)
|
||||
qpos_ref, qvel_ref, ctrl_ref, contact, contact_pos = load_data(
|
||||
config, config.data_path
|
||||
)
|
||||
if (
|
||||
config.contact_guidance
|
||||
and ctrl_ref.shape[1] != config.nu
|
||||
and qpos_ref.shape[1] >= config.nu
|
||||
):
|
||||
loguru.logger.info(
|
||||
"Using qpos as ctrl reference for contact guidance (ctrl dims: {} -> {}).",
|
||||
ctrl_ref.shape[1],
|
||||
config.nu,
|
||||
)
|
||||
ctrl_ref = qpos_ref[:, : config.nu]
|
||||
if config.contact_guidance and torch.all(contact <= 0):
|
||||
raise ValueError("contact_guidance is enabled, but contact mask is all zeros.")
|
||||
# hack: start from step 500
|
||||
# qpos_ref = qpos_ref[500:]
|
||||
# qvel_ref = qvel_ref[500:]
|
||||
# ctrl_ref = ctrl_ref[500:]
|
||||
# contact = contact[500:]
|
||||
# contact_pos = contact_pos[500:]
|
||||
ref_data = (qpos_ref, qvel_ref, ctrl_ref, contact, contact_pos)
|
||||
config.max_sim_steps = (
|
||||
config.max_sim_steps
|
||||
if config.max_sim_steps > 0
|
||||
else qpos_ref.shape[0] - config.horizon_steps - config.ctrl_steps
|
||||
)
|
||||
|
||||
# setup env with initial state from first sim qpos
|
||||
env = setup_env(config, ref_data)
|
||||
|
||||
# setup mujoco (for viewer only)
|
||||
mj_model = setup_mj_model(config)
|
||||
mj_data = mujoco.MjData(mj_model)
|
||||
mj_data_ref = mujoco.MjData(mj_model)
|
||||
mj_data.qpos[:] = qpos_ref[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = qvel_ref[0].detach().cpu().numpy()
|
||||
mj_data.ctrl[:] = ctrl_ref[0].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data)
|
||||
mj_data.time = 0.0
|
||||
_assert_object_actuator_gains_zero(env, config, "start")
|
||||
|
||||
# Initial-state stability sanity check (CPU MuJoCo only). Disabled when
|
||||
# config.sanity_check_seconds <= 0.0.
|
||||
if config.sanity_check_seconds > 0.0:
|
||||
_initial_state_sanity_check(
|
||||
mj_model,
|
||||
mj_data,
|
||||
qpos_ref,
|
||||
qvel_ref,
|
||||
ctrl_ref,
|
||||
config,
|
||||
save_video_dir=config.output_dir,
|
||||
)
|
||||
# Restore initial scene state after the check.
|
||||
mj_data.qpos[:] = qpos_ref[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = qvel_ref[0].detach().cpu().numpy()
|
||||
mj_data.ctrl[:] = ctrl_ref[0].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data)
|
||||
mj_data.time = 0.0
|
||||
|
||||
images = []
|
||||
object_trace_site_ids = []
|
||||
robot_trace_site_ids = []
|
||||
for sid in range(mj_model.nsite):
|
||||
name = mujoco.mj_id2name(mj_model, mujoco.mjtObj.mjOBJ_SITE, sid)
|
||||
if name is not None:
|
||||
if name.startswith("trace"):
|
||||
if "object" in name:
|
||||
object_trace_site_ids.append(sid)
|
||||
else:
|
||||
robot_trace_site_ids.append(sid)
|
||||
config.trace_site_ids = object_trace_site_ids + robot_trace_site_ids
|
||||
contact_guidance_enabled = (
|
||||
config.contact_guidance and len(config.object_actuator_ids) > 0
|
||||
)
|
||||
if config.contact_guidance and not contact_guidance_enabled:
|
||||
loguru.logger.warning(
|
||||
"contact_guidance is enabled but no object actuators were resolved."
|
||||
)
|
||||
contact_offset = 0
|
||||
if contact_guidance_enabled:
|
||||
config.contact_len = int(
|
||||
min(contact.shape[1], contact_pos.shape[1], len(config.contact_order))
|
||||
)
|
||||
if (
|
||||
config.contact_len != len(config.contact_order)
|
||||
or config.contact_len != contact.shape[1]
|
||||
):
|
||||
loguru.logger.warning(
|
||||
"Contact length mismatch (mask={}, pos={}, expected={}); truncating to {}.",
|
||||
contact.shape[1],
|
||||
contact_pos.shape[1],
|
||||
len(config.contact_order),
|
||||
config.contact_len,
|
||||
)
|
||||
config.contact_order = config.contact_order[: config.contact_len]
|
||||
config.hand_contact_site_ids = config.hand_contact_site_ids[
|
||||
: config.contact_len
|
||||
]
|
||||
contact_offset = max(contact.shape[1] - config.contact_len, 0)
|
||||
|
||||
# setup env params
|
||||
env_params_list = []
|
||||
if config.num_dr == 0:
|
||||
xy_offset_list = [0.0]
|
||||
pair_margin_list = [0.0]
|
||||
else:
|
||||
xy_offset_list = np.linspace(
|
||||
config.xy_offset_range[0], config.xy_offset_range[1], config.num_dr
|
||||
)
|
||||
pair_margin_list = np.linspace(
|
||||
config.pair_margin_range[0], config.pair_margin_range[1], config.num_dr
|
||||
)
|
||||
kp_schedule = []
|
||||
kd_schedule = []
|
||||
if contact_guidance_enabled and config.max_num_iterations > 0:
|
||||
actuator_names = config.object_actuator_names
|
||||
if not actuator_names:
|
||||
actuator_names = [
|
||||
mujoco.mj_id2name(mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, int(aid))
|
||||
for aid in config.object_actuator_ids
|
||||
]
|
||||
base_kp = np.array(
|
||||
[
|
||||
(
|
||||
config.init_rot_actuator_gain
|
||||
if ("_rot_" in (name or ""))
|
||||
else config.init_pos_actuator_gain
|
||||
)
|
||||
for name in actuator_names
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
base_kd = np.array(
|
||||
[
|
||||
(
|
||||
config.init_rot_actuator_bias
|
||||
if ("_rot_" in (name or ""))
|
||||
else config.init_pos_actuator_bias
|
||||
)
|
||||
for name in actuator_names
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
for i in range(config.max_num_iterations):
|
||||
decay = float(config.guidance_decay_ratio) ** i
|
||||
kp_i = base_kp * decay
|
||||
kd_i = base_kd * decay
|
||||
if i == config.max_num_iterations - 1:
|
||||
kp_i = np.zeros_like(base_kp, dtype=np.float32)
|
||||
kd_i = np.zeros_like(base_kd, dtype=np.float32)
|
||||
kp_schedule.append(kp_i)
|
||||
kd_schedule.append(kd_i)
|
||||
|
||||
for i in range(config.max_num_iterations):
|
||||
env_params = []
|
||||
for j in range(config.num_dr):
|
||||
params = {
|
||||
"xy_offset": xy_offset_list[j],
|
||||
"pair_margin": pair_margin_list[j],
|
||||
}
|
||||
if contact_guidance_enabled and kp_schedule:
|
||||
params["kp"] = kp_schedule[i]
|
||||
params["kd"] = kd_schedule[i]
|
||||
env_params.append(params)
|
||||
env_params_list.append(env_params)
|
||||
config.env_params_list = env_params_list
|
||||
_save_config_yaml(config)
|
||||
|
||||
# setup viewer and renderer
|
||||
run_viewer = setup_viewer(config, mj_model, mj_data)
|
||||
renderer = setup_renderer(config, mj_model)
|
||||
|
||||
# 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)
|
||||
base_noise_scale = config.noise_scale.clone()
|
||||
gibbs_enabled = config.gibbs_sampling and config.embodiment_type == "bimanual"
|
||||
if config.gibbs_sampling and not gibbs_enabled:
|
||||
loguru.logger.warning(
|
||||
"gibbs_sampling is enabled but embodiment_type is {}, disabling.",
|
||||
config.embodiment_type,
|
||||
)
|
||||
if gibbs_enabled:
|
||||
right_ids, left_ids = _get_bimanual_hand_indices(config)
|
||||
right_only_zero = left_ids
|
||||
left_only_zero = right_ids
|
||||
|
||||
# initial controls
|
||||
ctrls = ctrl_ref[: config.horizon_steps]
|
||||
# buffers for saving info and trajectory
|
||||
info_list = []
|
||||
|
||||
# run viewer + control loop
|
||||
t_start = time.perf_counter()
|
||||
with run_viewer() as viewer:
|
||||
while viewer.is_running():
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# optimize using future reference window at control-rate (+1 lookahead)
|
||||
sim_step = int(np.round(mj_data.time / config.sim_dt))
|
||||
ref_slice = get_slice(
|
||||
ref_data, sim_step + 1, sim_step + config.horizon_steps + 1
|
||||
)
|
||||
ctrls_for_opt = ctrls
|
||||
if contact_guidance_enabled and config.contact_len > 0:
|
||||
contact_mask_step = contact[sim_step][
|
||||
contact_offset : contact_offset + config.contact_len
|
||||
]
|
||||
contact_pos_ref_step = contact_pos[sim_step]
|
||||
site_xpos = wp.to_torch(env.data_wp.site_xpos)[0]
|
||||
|
||||
right_delta = compute_contact_point_delta(
|
||||
contact_mask_step,
|
||||
contact_pos_ref_step,
|
||||
site_xpos,
|
||||
config.hand_contact_site_ids,
|
||||
config.right_contact_indices,
|
||||
)
|
||||
left_delta = compute_contact_point_delta(
|
||||
contact_mask_step,
|
||||
contact_pos_ref_step,
|
||||
site_xpos,
|
||||
config.hand_contact_site_ids,
|
||||
config.left_contact_indices,
|
||||
)
|
||||
if (
|
||||
right_delta is not None
|
||||
and config.right_pos_ctrl_ids
|
||||
and sim_step + ctrls.shape[0] <= ctrl_ref.shape[0]
|
||||
):
|
||||
ctrls_for_opt = ctrls_for_opt.clone()
|
||||
ref_ctrl_slice = ctrl_ref[sim_step : sim_step + ctrls.shape[0]]
|
||||
ctrls_for_opt[:, config.right_pos_ctrl_ids] = ref_ctrl_slice[
|
||||
:, config.right_pos_ctrl_ids
|
||||
] + torch.clip(right_delta, -0.01, 0.01)
|
||||
if (
|
||||
left_delta is not None
|
||||
and config.left_pos_ctrl_ids
|
||||
and sim_step + ctrls.shape[0] <= ctrl_ref.shape[0]
|
||||
):
|
||||
if ctrls_for_opt is ctrls:
|
||||
ctrls_for_opt = ctrls_for_opt.clone()
|
||||
ref_ctrl_slice = ctrl_ref[sim_step : sim_step + ctrls.shape[0]]
|
||||
ctrls_for_opt[:, config.left_pos_ctrl_ids] = ref_ctrl_slice[
|
||||
:, config.left_pos_ctrl_ids
|
||||
] + torch.clip(left_delta, -0.01, 0.01)
|
||||
if gibbs_enabled:
|
||||
config.noise_scale = _apply_noise_mask(
|
||||
base_noise_scale, right_only_zero
|
||||
)
|
||||
ctrls, infos = optimize(config, env, ctrls_for_opt, ref_slice)
|
||||
config.noise_scale = _apply_noise_mask(base_noise_scale, left_only_zero)
|
||||
ctrls, infos = optimize(config, env, ctrls, ref_slice)
|
||||
config.noise_scale = base_noise_scale
|
||||
else:
|
||||
config.noise_scale = base_noise_scale
|
||||
ctrls, infos = optimize(config, env, ctrls_for_opt, ref_slice)
|
||||
|
||||
# Compute trace_ref from reference qpos over the horizon
|
||||
if len(config.trace_site_ids) > 0:
|
||||
trace_ref = []
|
||||
qpos_ref_horizon = ref_slice[0]
|
||||
for h in range(config.horizon_steps):
|
||||
mj_data_ref.qpos[:] = qpos_ref_horizon[h].detach().cpu().numpy()
|
||||
mujoco.mj_kinematics(mj_model, mj_data_ref)
|
||||
site_xpos = np.array(
|
||||
[mj_data_ref.site_xpos[sid] for sid in config.trace_site_ids]
|
||||
)
|
||||
trace_ref.append(site_xpos)
|
||||
# (H, K, 3) -> (1, 1, H, K, 3) to match trace_sample shape
|
||||
trace_ref_np = np.stack(trace_ref, axis=0)[None, None, :, :, :]
|
||||
infos["trace_ref"] = trace_ref_np
|
||||
|
||||
# step environment for ctrl_steps
|
||||
step_info = {"qpos": [], "qvel": [], "time": [], "ctrl": []}
|
||||
for i in range(config.ctrl_steps):
|
||||
ctrl_step = ctrls[i]
|
||||
|
||||
# option 1: use mujoco step
|
||||
# mj_data.ctrl[:] = ctrls[i].detach().cpu().numpy()
|
||||
# mujoco.mj_step(mj_model, mj_data)
|
||||
# option 2: use warp step
|
||||
step_env(config, env, ctrl_step)
|
||||
mj_data.qpos[:] = get_qpos(config, env)[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = get_qvel(config, env)[0].detach().cpu().numpy()
|
||||
mj_data.ctrl[:] = ctrl_step.detach().cpu().numpy()
|
||||
mj_data.time += config.sim_dt
|
||||
if config.save_video and renderer is not None:
|
||||
if i % int(np.round(config.render_dt / config.sim_dt)) == 0:
|
||||
mj_data_ref.qpos[:] = (
|
||||
qpos_ref[sim_step + i].detach().cpu().numpy()
|
||||
)
|
||||
image = render_image(
|
||||
config, renderer, mj_model, mj_data, mj_data_ref
|
||||
)
|
||||
images.append(image)
|
||||
if "rerun" in config.viewer or "viser" in config.viewer:
|
||||
mj_data_ref.qpos[:] = qpos_ref[sim_step + i].detach().cpu().numpy()
|
||||
mujoco.mj_kinematics(mj_model, mj_data_ref)
|
||||
log_frame(
|
||||
mj_data,
|
||||
sim_time=mj_data.time,
|
||||
viewer_body_entity_and_ids=config.viewer_body_entity_and_ids,
|
||||
data_ref=mj_data_ref,
|
||||
)
|
||||
step_info["qpos"].append(mj_data.qpos.copy())
|
||||
step_info["qvel"].append(mj_data.qvel.copy())
|
||||
step_info["time"].append(mj_data.time)
|
||||
step_info["ctrl"].append(mj_data.ctrl.copy())
|
||||
for k in step_info:
|
||||
step_info[k] = np.stack(step_info[k], axis=0)
|
||||
infos.update(step_info)
|
||||
# sync env state
|
||||
sync_env(config, env, mj_data)
|
||||
|
||||
# receding horizon update
|
||||
sim_step = int(np.round(mj_data.time / config.sim_dt))
|
||||
prev_ctrl = ctrls[config.ctrl_steps :]
|
||||
new_ctrl = ctrl_ref[
|
||||
sim_step + prev_ctrl.shape[0] : sim_step
|
||||
+ prev_ctrl.shape[0]
|
||||
+ config.ctrl_steps
|
||||
]
|
||||
ctrls = torch.cat([prev_ctrl, new_ctrl], dim=0)
|
||||
|
||||
# sync viewer state and render
|
||||
mj_data.qpos[:] = get_qpos(config, env)[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = get_qvel(config, env)[0].detach().cpu().numpy()
|
||||
mj_data_ref.qpos[:] = qpos_ref[sim_step].detach().cpu().numpy()
|
||||
update_viewer(config, viewer, mj_model, mj_data, mj_data_ref, 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"Total time: {t_end - t_start:.4f}s")
|
||||
|
||||
# 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_mjwp{'_act' if config.contact_guidance else ''}.npz",
|
||||
**info_aggregated,
|
||||
)
|
||||
loguru.logger.info(
|
||||
f"Saved info to {config.output_dir}/trajectory_mjwp{'_act' if config.contact_guidance else ''}.npz"
|
||||
)
|
||||
|
||||
# save video
|
||||
if config.save_video and len(images) > 0:
|
||||
video_path = f"{config.output_dir}/visualization_mjwp{'_act' if config.contact_guidance else ''}.mp4"
|
||||
imageio.mimsave(
|
||||
video_path,
|
||||
images,
|
||||
fps=int(1 / config.render_dt),
|
||||
)
|
||||
loguru.logger.info(f"Saved video to {video_path}")
|
||||
|
||||
errors = None
|
||||
if info_list:
|
||||
qpos_traj = np.concatenate([info["qpos"] for info in info_list], axis=0)
|
||||
qpos_ref_np = qpos_ref[: qpos_traj.shape[0]].detach().cpu().numpy()
|
||||
data_type = "mjwp_act" if config.contact_guidance else "mjwp"
|
||||
errors = compute_object_tracking_error(
|
||||
qpos_traj, qpos_ref_np, config.embodiment_type, data_type
|
||||
)
|
||||
loguru.logger.info(
|
||||
"Final object tracking error: pos={:.4f}, quat={:.4f}",
|
||||
errors["obj_pos_err"],
|
||||
errors["obj_quat_err"],
|
||||
)
|
||||
|
||||
_assert_object_actuator_gains_zero(env, config, "end")
|
||||
|
||||
if "viser" in config.viewer and config.wait_on_finish:
|
||||
loguru.logger.info(
|
||||
"Optimization complete! Keeping Viser server alive. Press Ctrl+C to exit."
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="config", config_name="default")
|
||||
def run_main(cfg: DictConfig) -> None:
|
||||
"""Entry point for Hydra configuration runner."""
|
||||
# Convert DictConfig to Config dataclass, handling special fields
|
||||
config_dict = dict(cfg)
|
||||
|
||||
# Optionally load a saved config YAML and merge; CLI overrides take priority.
|
||||
load_config_path = config_dict.get("load_config_path", "")
|
||||
if load_config_path:
|
||||
loaded_config = load_config_yaml(load_config_path)
|
||||
cli_overrides = _extract_cli_overrides(cfg)
|
||||
config_dict = {**loaded_config, **cli_overrides}
|
||||
else:
|
||||
config_dict = filter_config_fields(config_dict)
|
||||
|
||||
# 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()
|
||||
+834
@@ -0,0 +1,834 @@
|
||||
# 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.
|
||||
|
||||
"""Deterministic optimization variant of run_mjwp.py.
|
||||
|
||||
Instead of re-rolling out the optimized controls, this script replays the
|
||||
pre-computed best trajectory states and restores the full simulation state
|
||||
from the optimizer's snapshot.
|
||||
|
||||
Retry/revert is modeled as an explicit state machine — see `_StepState`
|
||||
and the transition table at the top of `main()`.
|
||||
|
||||
Author: Chaoyi Pan
|
||||
Date: 2025-04-16
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field, fields
|
||||
from pathlib import Path
|
||||
|
||||
import hydra
|
||||
import imageio
|
||||
import loguru
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from spider.config import (
|
||||
Config,
|
||||
filter_config_fields,
|
||||
load_config_yaml,
|
||||
process_config,
|
||||
)
|
||||
from spider.interp import get_slice
|
||||
from spider.io import load_data
|
||||
from spider.optimizers.sampling_fast import (
|
||||
_restore_best_state_to_env,
|
||||
make_optimize_fn_fast,
|
||||
make_optimize_once_fn_fast,
|
||||
make_rollout_fn_fast,
|
||||
)
|
||||
from spider.postprocess.get_success_rate import compute_object_tracking_error
|
||||
from spider.simulators.mjwp import (
|
||||
_initial_state_sanity_check,
|
||||
copy_sample_state,
|
||||
get_qpos,
|
||||
get_qvel,
|
||||
get_reward,
|
||||
get_terminal_reward,
|
||||
get_terminate,
|
||||
get_trace,
|
||||
load_env_params,
|
||||
load_state,
|
||||
save_env_params,
|
||||
save_state,
|
||||
setup_env,
|
||||
setup_mj_model,
|
||||
step_env,
|
||||
sync_env_mujoco,
|
||||
)
|
||||
from spider.viewers import (
|
||||
log_frame,
|
||||
render_image,
|
||||
setup_renderer,
|
||||
setup_viewer,
|
||||
update_viewer,
|
||||
)
|
||||
|
||||
_CONFIG_SKIP_FIELDS = {
|
||||
"noise_scale",
|
||||
"env_params_list",
|
||||
"viewer_body_entity_and_ids",
|
||||
}
|
||||
|
||||
|
||||
def _parse_override_tokens(tokens: list[str]) -> dict:
|
||||
allowed = {field.name for field in fields(Config)}
|
||||
override_dict: dict = {}
|
||||
for item in tokens:
|
||||
if "=" not in item:
|
||||
continue
|
||||
key, value = item.split("=", 1)
|
||||
key = key.lstrip("+")
|
||||
if key not in allowed:
|
||||
continue
|
||||
parsed = OmegaConf.to_container(
|
||||
OmegaConf.from_dotlist([f"{key}={value}"]), resolve=True
|
||||
)
|
||||
if isinstance(parsed, dict) and key in parsed:
|
||||
override_dict[key] = parsed[key]
|
||||
return override_dict
|
||||
|
||||
|
||||
def _extract_cli_overrides(cfg: DictConfig) -> dict:
|
||||
"""Extract CLI overrides so they can be applied on top of a loaded config."""
|
||||
overrides = OmegaConf.select(cfg, "hydra.overrides.task") or []
|
||||
override_dict = _parse_override_tokens(overrides)
|
||||
if override_dict:
|
||||
return override_dict
|
||||
return _parse_override_tokens(sys.argv[1:])
|
||||
|
||||
|
||||
def _normalize_yaml_value(value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.detach().cpu().tolist()
|
||||
if isinstance(value, np.ndarray):
|
||||
return value.tolist()
|
||||
if isinstance(value, np.generic):
|
||||
return value.item()
|
||||
if isinstance(value, tuple):
|
||||
return list(value)
|
||||
return value
|
||||
|
||||
|
||||
def _save_config_yaml(config: Config) -> None:
|
||||
if not config.save_config:
|
||||
return
|
||||
config_dict = {}
|
||||
for fld in fields(config):
|
||||
if fld.name in _CONFIG_SKIP_FIELDS:
|
||||
continue
|
||||
config_dict[fld.name] = _normalize_yaml_value(getattr(config, fld.name))
|
||||
output_path = Path(config.output_dir) / "config.yaml"
|
||||
OmegaConf.save(config=OmegaConf.create(config_dict), f=str(output_path))
|
||||
loguru.logger.info(f"Saved config to {output_path}")
|
||||
|
||||
|
||||
# --- Progressive retry boost ------------------------------------------------
|
||||
#
|
||||
# On each retry/revert pop, boost_level increments. Iterations double per
|
||||
# level, noise scales by sqrt(2) per level — both fixed-base (relative to
|
||||
# the captured originals), so level=2 is orig*4 iterations, orig*2 noise.
|
||||
|
||||
_BOOST_NOISE_FACTOR = math.sqrt(2.0)
|
||||
|
||||
|
||||
def _apply_boost(
|
||||
config: Config,
|
||||
boost_level: int,
|
||||
orig_max_iters: int,
|
||||
orig_env_params_list: list,
|
||||
orig_noise_scale: torch.Tensor,
|
||||
) -> None:
|
||||
iters = orig_max_iters * (2**boost_level)
|
||||
config.max_num_iterations = iters
|
||||
config.beta_traj = config.final_noise_scale ** (1.0 / iters)
|
||||
extended = list(orig_env_params_list)
|
||||
while len(extended) < iters:
|
||||
extended.append(extended[-1])
|
||||
config.env_params_list = extended
|
||||
config.noise_scale = orig_noise_scale * (_BOOST_NOISE_FACTOR**boost_level)
|
||||
|
||||
|
||||
# --- Retry/revert state machine ---------------------------------------------
|
||||
#
|
||||
# States and transitions (event -> action):
|
||||
#
|
||||
# OPTIMIZE succeed -> OPTIMIZE
|
||||
# fail, local<cap -> LOCAL_RETRY
|
||||
# fail, local>=cap, history -> REVERT (first)
|
||||
# fail, local>=cap, no history -> DONE
|
||||
#
|
||||
# LOCAL_RETRY (transient) reload curr_cp, boost -> (unchanged state)
|
||||
#
|
||||
# REVERT (first, from OPTIMIZE) pop once, save anchor -> RECOVERING
|
||||
#
|
||||
# RECOVERING succeed, sim_step > anchor -> OPTIMIZE (full recovery)
|
||||
# succeed, sim_step <= anchor -> RECOVERING (no push)
|
||||
# fail, local<cap -> LOCAL_RETRY
|
||||
# fail, local>=cap, attempts<max -> RE-REVERT (same anchor)
|
||||
# fail, local>=cap, attempts>=max -> DONE
|
||||
#
|
||||
# "Recovery anchor" = sim_step right after the first REVERT pop. On subsequent
|
||||
# failures during recovery, we reload the *same* anchor checkpoint instead of
|
||||
# popping deeper into history. Recovery is declared when a step succeeds at
|
||||
# sim_step > anchor.
|
||||
|
||||
|
||||
class _StepState:
|
||||
OPTIMIZE = "OPTIMIZE"
|
||||
RECOVERING = "RECOVERING"
|
||||
DONE = "DONE"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LoopCtx:
|
||||
state: str = _StepState.OPTIMIZE
|
||||
history: collections.deque = field(default_factory=collections.deque)
|
||||
local_retries: int = 0
|
||||
boost_level: int = 0
|
||||
revert_anchor_step: int | None = None
|
||||
revert_anchor_checkpoint: dict | None = None
|
||||
revert_attempts_from_anchor: int = 0
|
||||
|
||||
|
||||
def main(config: Config):
|
||||
"""Run SPIDER using MuJoCo Warp backend with deterministic best-trajectory replay."""
|
||||
config = process_config(config)
|
||||
|
||||
if config.contact_guidance:
|
||||
raise NotImplementedError(
|
||||
"contact_guidance is not supported in run_mjwp_fast.py. "
|
||||
"Use examples/run_mjwp.py for contact-guided runs."
|
||||
)
|
||||
if config.gibbs_sampling:
|
||||
raise NotImplementedError(
|
||||
"gibbs_sampling is not supported in run_mjwp_fast.py. "
|
||||
"Use examples/run_mjwp.py for Gibbs sampling."
|
||||
)
|
||||
|
||||
qpos_ref, qvel_ref, ctrl_ref, contact, contact_pos = load_data(
|
||||
config, config.data_path
|
||||
)
|
||||
|
||||
ref_data = (qpos_ref, qvel_ref, ctrl_ref, contact, contact_pos)
|
||||
config.max_sim_steps = (
|
||||
config.max_sim_steps
|
||||
if config.max_sim_steps > 0
|
||||
else qpos_ref.shape[0] - config.horizon_steps - config.ctrl_steps
|
||||
)
|
||||
|
||||
env = setup_env(config, ref_data)
|
||||
|
||||
mj_model = setup_mj_model(config)
|
||||
mj_data = mujoco.MjData(mj_model)
|
||||
mj_data_ref = mujoco.MjData(mj_model)
|
||||
mj_data.qpos[:] = qpos_ref[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = qvel_ref[0].detach().cpu().numpy()
|
||||
mj_data.ctrl[:] = ctrl_ref[0].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data)
|
||||
mj_data.time = 0.0
|
||||
|
||||
# Initial-state stability sanity check (CPU MuJoCo only). Disabled when
|
||||
# config.sanity_check_seconds <= 0.0.
|
||||
if config.sanity_check_seconds > 0.0:
|
||||
_initial_state_sanity_check(
|
||||
mj_model,
|
||||
mj_data,
|
||||
qpos_ref,
|
||||
qvel_ref,
|
||||
ctrl_ref,
|
||||
config,
|
||||
save_video_dir=config.output_dir,
|
||||
)
|
||||
# Restore initial scene state after the check.
|
||||
mj_data.qpos[:] = qpos_ref[0].detach().cpu().numpy()
|
||||
mj_data.qvel[:] = qvel_ref[0].detach().cpu().numpy()
|
||||
mj_data.ctrl[:] = ctrl_ref[0].detach().cpu().numpy()
|
||||
mujoco.mj_step(mj_model, mj_data)
|
||||
mj_data.time = 0.0
|
||||
|
||||
images = []
|
||||
object_trace_site_ids = []
|
||||
robot_trace_site_ids = []
|
||||
for sid in range(mj_model.nsite):
|
||||
name = mujoco.mj_id2name(mj_model, mujoco.mjtObj.mjOBJ_SITE, sid)
|
||||
if name is not None and name.startswith("trace"):
|
||||
if "object" in name:
|
||||
object_trace_site_ids.append(sid)
|
||||
else:
|
||||
robot_trace_site_ids.append(sid)
|
||||
config.trace_site_ids = object_trace_site_ids + robot_trace_site_ids
|
||||
|
||||
# setup env params (no DR-schedule dependency on contact guidance)
|
||||
if config.num_dr == 0:
|
||||
xy_offset_list = [0.0]
|
||||
pair_margin_list = [0.0]
|
||||
else:
|
||||
xy_offset_list = np.linspace(
|
||||
config.xy_offset_range[0], config.xy_offset_range[1], config.num_dr
|
||||
)
|
||||
pair_margin_list = np.linspace(
|
||||
config.pair_margin_range[0], config.pair_margin_range[1], config.num_dr
|
||||
)
|
||||
env_params_list = []
|
||||
for _ in range(config.max_num_iterations):
|
||||
env_params = [
|
||||
{"xy_offset": xy_offset_list[j], "pair_margin": pair_margin_list[j]}
|
||||
for j in range(config.num_dr)
|
||||
]
|
||||
env_params_list.append(env_params)
|
||||
config.env_params_list = env_params_list
|
||||
_save_config_yaml(config)
|
||||
|
||||
# setup viewer and renderer
|
||||
run_viewer = setup_viewer(config, mj_model, mj_data)
|
||||
renderer = setup_renderer(config, mj_model)
|
||||
|
||||
# If viser is enabled and the user wants the recording saved, start
|
||||
# recording NOW so every per-frame transform update from log_frame is
|
||||
# captured to a serializable timeline (a single .viser file).
|
||||
_viser_recording_active = False
|
||||
if "viser" in (config.viewer or "").lower() and config.save_viser:
|
||||
try:
|
||||
from spider.viewers.viser_viewer import start_recording
|
||||
|
||||
start_recording(dt=float(config.render_dt))
|
||||
_viser_recording_active = True
|
||||
except Exception as e:
|
||||
loguru.logger.warning(f"Failed to start viser recording: {e}")
|
||||
|
||||
# Register an atexit + SIGTERM handler so timeout-killed runs still
|
||||
# flush whatever optimization video / viser timeline was captured.
|
||||
# Without this, mp4 + .viser are only written if main() reaches its
|
||||
# end-of-loop save block, which fails-to-converge runs never do.
|
||||
import atexit
|
||||
import os as _os
|
||||
import signal as _signal
|
||||
|
||||
_flushed = {"done": False}
|
||||
|
||||
def _emergency_flush() -> None:
|
||||
if _flushed["done"]:
|
||||
return
|
||||
_flushed["done"] = True
|
||||
# Save partial mp4 if any frames are queued
|
||||
try:
|
||||
if config.save_video and len(images) > 0:
|
||||
video_path = (
|
||||
f"{config.output_dir}/visualization_mjwp_fast.mp4"
|
||||
)
|
||||
imageio.mimsave(
|
||||
video_path,
|
||||
images,
|
||||
fps=int(1 / config.render_dt),
|
||||
)
|
||||
loguru.logger.warning(
|
||||
f"[emergency flush] Saved partial video ({len(images)} frames) to {video_path}"
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
loguru.logger.warning(f"[emergency flush] mp4 save failed: {e}")
|
||||
# Save viser timeline if recording was started
|
||||
try:
|
||||
if _viser_recording_active:
|
||||
from spider.viewers.viser_viewer import end_recording_and_save
|
||||
|
||||
viser_path = (
|
||||
f"{config.output_dir}/visualization_mjwp_fast.viser"
|
||||
)
|
||||
if end_recording_and_save(viser_path):
|
||||
loguru.logger.warning(
|
||||
f"[emergency flush] Saved partial viser timeline to {viser_path}"
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
loguru.logger.warning(f"[emergency flush] viser save failed: {e}")
|
||||
|
||||
atexit.register(_emergency_flush)
|
||||
|
||||
def _on_sigterm(_sig, _frame) -> None:
|
||||
loguru.logger.warning("Received SIGTERM; flushing artifacts.")
|
||||
_emergency_flush()
|
||||
# Re-raise default handler so the process actually exits.
|
||||
_signal.signal(_signal.SIGTERM, _signal.SIG_DFL)
|
||||
_os.kill(_os.getpid(), _signal.SIGTERM)
|
||||
|
||||
try:
|
||||
_signal.signal(_signal.SIGTERM, _on_sigterm)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
# setup optimizer — use fast variants
|
||||
rollout = make_rollout_fn_fast(
|
||||
step_env,
|
||||
save_state,
|
||||
load_state,
|
||||
get_reward,
|
||||
get_terminal_reward,
|
||||
get_terminate,
|
||||
get_trace,
|
||||
save_env_params,
|
||||
load_env_params,
|
||||
copy_sample_state,
|
||||
get_qpos,
|
||||
get_qvel,
|
||||
)
|
||||
optimize_once = make_optimize_once_fn_fast(rollout)
|
||||
optimize = make_optimize_fn_fast(optimize_once)
|
||||
|
||||
ctrls = ctrl_ref[: config.horizon_steps]
|
||||
info_list = []
|
||||
|
||||
orig_max_iters = config.max_num_iterations
|
||||
orig_env_params_list = list(config.env_params_list)
|
||||
orig_noise_scale = config.noise_scale.clone()
|
||||
|
||||
last_snap = None
|
||||
last_snap_idx = None
|
||||
|
||||
def _save_checkpoint():
|
||||
return {
|
||||
"qpos": mj_data.qpos.copy(),
|
||||
"qvel": mj_data.qvel.copy(),
|
||||
"ctrl": mj_data.ctrl.copy(),
|
||||
"time": mj_data.time,
|
||||
"ctrls": ctrls.clone(),
|
||||
"env_snap": last_snap,
|
||||
"env_idx": last_snap_idx,
|
||||
"images_len": len(images),
|
||||
"info_list_len": len(info_list),
|
||||
}
|
||||
|
||||
def _load_checkpoint(cp):
|
||||
nonlocal ctrls, last_snap, last_snap_idx
|
||||
mj_data.qpos[:] = cp["qpos"]
|
||||
mj_data.qvel[:] = cp["qvel"]
|
||||
mj_data.ctrl[:] = cp["ctrl"]
|
||||
mj_data.time = cp["time"]
|
||||
mujoco.mj_forward(mj_model, mj_data)
|
||||
ctrls = cp["ctrls"].clone()
|
||||
last_snap = cp["env_snap"]
|
||||
last_snap_idx = cp["env_idx"]
|
||||
if last_snap is not None:
|
||||
_restore_best_state_to_env(env, last_snap, last_snap_idx)
|
||||
else:
|
||||
sync_env_mujoco(config, env, mj_data)
|
||||
del images[cp["images_len"] :]
|
||||
del info_list[cp["info_list_len"] :]
|
||||
|
||||
def _run_optimize():
|
||||
nonlocal ctrls
|
||||
s = int(np.round(mj_data.time / config.sim_dt))
|
||||
rs = get_slice(ref_data, s + 1, s + config.horizon_steps + 1)
|
||||
ctrls, inf = optimize(config, env, ctrls, rs)
|
||||
if len(config.trace_site_ids) > 0:
|
||||
tr = []
|
||||
qrh = rs[0]
|
||||
for h in range(config.horizon_steps):
|
||||
mj_data_ref.qpos[:] = qrh[h].detach().cpu().numpy()
|
||||
mujoco.mj_kinematics(mj_model, mj_data_ref)
|
||||
tr.append(
|
||||
np.array(
|
||||
[mj_data_ref.site_xpos[sid] for sid in config.trace_site_ids]
|
||||
)
|
||||
)
|
||||
inf["trace_ref"] = np.stack(tr, axis=0)[None, None, :, :, :]
|
||||
return s, rs, inf
|
||||
|
||||
def _replay_and_advance(sim_step_val, inf):
|
||||
nonlocal ctrls
|
||||
bqp = torch.from_numpy(inf["best_qpos"])
|
||||
bqv = torch.from_numpy(inf["best_qvel"])
|
||||
si = {"qpos": [], "qvel": [], "time": [], "ctrl": []}
|
||||
for i in range(config.ctrl_steps):
|
||||
mj_data.qpos[:] = bqp[i].numpy()
|
||||
mj_data.qvel[:] = bqv[i].numpy()
|
||||
mj_data.ctrl[:] = ctrls[i].detach().cpu().numpy()
|
||||
mj_data.time += config.sim_dt
|
||||
mujoco.mj_forward(mj_model, mj_data)
|
||||
if (
|
||||
config.save_video
|
||||
and renderer is not None
|
||||
and i % int(np.round(config.render_dt / config.sim_dt)) == 0
|
||||
):
|
||||
mj_data_ref.qpos[:] = qpos_ref[sim_step_val + i].detach().cpu().numpy()
|
||||
img = render_image(config, renderer, mj_model, mj_data, mj_data_ref)
|
||||
images.append(img)
|
||||
if "rerun" in config.viewer or "viser" in config.viewer:
|
||||
mj_data_ref.qpos[:] = qpos_ref[sim_step_val + i].detach().cpu().numpy()
|
||||
mujoco.mj_kinematics(mj_model, mj_data_ref)
|
||||
log_frame(
|
||||
mj_data,
|
||||
sim_time=mj_data.time,
|
||||
viewer_body_entity_and_ids=config.viewer_body_entity_and_ids,
|
||||
data_ref=mj_data_ref,
|
||||
)
|
||||
# Tick the optimization-process recording so this frame is
|
||||
# placed at its own timestamp on the timeline.
|
||||
if "viser" in config.viewer and config.save_viser:
|
||||
try:
|
||||
from spider.viewers.viser_viewer import tick_recording
|
||||
|
||||
tick_recording()
|
||||
except Exception:
|
||||
pass
|
||||
si["qpos"].append(mj_data.qpos.copy())
|
||||
si["qvel"].append(mj_data.qvel.copy())
|
||||
si["time"].append(mj_data.time)
|
||||
si["ctrl"].append(mj_data.ctrl.copy())
|
||||
for k in si:
|
||||
si[k] = np.stack(si[k], axis=0)
|
||||
inf.update(si)
|
||||
_restore_best_state_to_env(env, inf["state_snapshot"], inf["best_idx"])
|
||||
new_sim_step = int(np.round(mj_data.time / config.sim_dt))
|
||||
prev_c = ctrls[config.ctrl_steps :]
|
||||
new_c = ctrl_ref[
|
||||
new_sim_step + prev_c.shape[0] : new_sim_step
|
||||
+ prev_c.shape[0]
|
||||
+ config.ctrl_steps
|
||||
]
|
||||
ctrls = torch.cat([prev_c, new_c], dim=0)
|
||||
mj_data_ref.qpos[:] = qpos_ref[new_sim_step].detach().cpu().numpy()
|
||||
update_viewer(config, viewer, mj_model, mj_data, mj_data_ref, inf)
|
||||
info_list.append(
|
||||
{
|
||||
k: v
|
||||
for k, v in inf.items()
|
||||
if k not in ["trace_sample", "state_snapshot"]
|
||||
}
|
||||
)
|
||||
return new_sim_step
|
||||
|
||||
def _full_reset():
|
||||
ctx.boost_level = 0
|
||||
ctx.revert_anchor_step = None
|
||||
ctx.revert_anchor_checkpoint = None
|
||||
ctx.revert_attempts_from_anchor = 0
|
||||
ctx.local_retries = 0
|
||||
_apply_boost(config, 0, orig_max_iters, orig_env_params_list, orig_noise_scale)
|
||||
|
||||
def _truncate_to_latest_commit():
|
||||
if ctx.history:
|
||||
anchor_cp = ctx.history[-1]
|
||||
del images[anchor_cp["images_len"] :]
|
||||
del info_list[anchor_cp["info_list_len"] :]
|
||||
else:
|
||||
del images[:]
|
||||
del info_list[:]
|
||||
|
||||
ctx = _LoopCtx()
|
||||
# sim_step is defined in the loop; keep a module-visible default for the
|
||||
# `sim_step >= max_sim_steps` terminal check.
|
||||
sim_step = 0
|
||||
|
||||
t_start = time.perf_counter()
|
||||
with run_viewer() as viewer:
|
||||
while ctx.state != _StepState.DONE and viewer.is_running():
|
||||
t0 = time.perf_counter()
|
||||
curr_checkpoint = _save_checkpoint()
|
||||
sim_step, ref_slice, infos = _run_optimize()
|
||||
|
||||
# --- Per-step tracking error for diagnostics ---
|
||||
bqp_diag = infos.get("best_qpos")
|
||||
err_info = {"max_pos_error": float("nan"), "max_rot_error": float("nan")}
|
||||
if bqp_diag is not None:
|
||||
from spider.optimizers.sampling_fast import check_tracking_error
|
||||
|
||||
bqp_t = (
|
||||
torch.from_numpy(bqp_diag)
|
||||
if isinstance(bqp_diag, np.ndarray)
|
||||
else bqp_diag
|
||||
).to(config.device)
|
||||
# Same alignment as in sampling_fast.optimize: recorded_qpos[i]
|
||||
# lives at ref index sim_step+1+i = ref_slice[0][i] (NOT [i+1]).
|
||||
avail = min(config.ctrl_steps, ref_slice[0].shape[0])
|
||||
ref_qpos_diag = ref_slice[0][:avail].to(config.device)
|
||||
_, err_info = check_tracking_error(config, bqp_t[:avail], ref_qpos_diag)
|
||||
infos["step_pos_error"] = np.array([err_info["max_pos_error"]])
|
||||
infos["step_rot_error"] = np.array([err_info["max_rot_error"]])
|
||||
|
||||
succeeded = bool(infos["succeeded"][0])
|
||||
|
||||
if succeeded:
|
||||
# --- SUCCESS ---
|
||||
full_recovery = (
|
||||
ctx.state == _StepState.RECOVERING
|
||||
and ctx.revert_anchor_step is not None
|
||||
and sim_step > ctx.revert_anchor_step
|
||||
)
|
||||
if full_recovery:
|
||||
loguru.logger.info(
|
||||
"Recovered after {} revert attempts at sim_step={}.",
|
||||
ctx.revert_attempts_from_anchor,
|
||||
sim_step,
|
||||
)
|
||||
_full_reset()
|
||||
ctx.state = _StepState.OPTIMIZE
|
||||
|
||||
ctx.local_retries = 0
|
||||
if ctx.state == _StepState.OPTIMIZE:
|
||||
ctx.history.append(curr_checkpoint)
|
||||
last_snap = infos["state_snapshot"]
|
||||
last_snap_idx = infos["best_idx"]
|
||||
sim_step = _replay_and_advance(sim_step, infos)
|
||||
|
||||
elif ctx.local_retries < config.max_local_retries:
|
||||
# --- LOCAL_RETRY ---
|
||||
ctx.local_retries += 1
|
||||
loguru.logger.warning(
|
||||
"Local retry {}/{} at sim_step={} (pos={:.4f}, rot={:.4f}).",
|
||||
ctx.local_retries,
|
||||
config.max_local_retries,
|
||||
sim_step,
|
||||
err_info["max_pos_error"],
|
||||
err_info["max_rot_error"],
|
||||
)
|
||||
_load_checkpoint(curr_checkpoint)
|
||||
ctx.boost_level += 1
|
||||
_apply_boost(
|
||||
config,
|
||||
ctx.boost_level,
|
||||
orig_max_iters,
|
||||
orig_env_params_list,
|
||||
orig_noise_scale,
|
||||
)
|
||||
continue
|
||||
|
||||
elif (
|
||||
ctx.state == _StepState.RECOVERING
|
||||
and ctx.revert_anchor_checkpoint is not None
|
||||
):
|
||||
# --- RE-REVERT to same anchor ---
|
||||
ctx.revert_attempts_from_anchor += 1
|
||||
if ctx.revert_attempts_from_anchor > config.max_revert_depth:
|
||||
loguru.logger.error(
|
||||
"Exhausted {} revert attempts from anchor sim_step={}.",
|
||||
config.max_revert_depth,
|
||||
ctx.revert_anchor_step,
|
||||
)
|
||||
_apply_boost(
|
||||
config,
|
||||
0,
|
||||
orig_max_iters,
|
||||
orig_env_params_list,
|
||||
orig_noise_scale,
|
||||
)
|
||||
_truncate_to_latest_commit()
|
||||
sim_step = _replay_and_advance(sim_step, infos)
|
||||
ctx.state = _StepState.DONE
|
||||
break
|
||||
ctx.local_retries = 0
|
||||
_load_checkpoint(ctx.revert_anchor_checkpoint)
|
||||
# Reset boost so iters/noise don't compound across anchors —
|
||||
# each new anchor starts a fresh local-retry sweep at level 0.
|
||||
ctx.boost_level = 0
|
||||
_apply_boost(
|
||||
config,
|
||||
ctx.boost_level,
|
||||
orig_max_iters,
|
||||
orig_env_params_list,
|
||||
orig_noise_scale,
|
||||
)
|
||||
loguru.logger.warning(
|
||||
"Re-revert attempt {}/{} back to anchor sim_step={}.",
|
||||
ctx.revert_attempts_from_anchor,
|
||||
config.max_revert_depth,
|
||||
ctx.revert_anchor_step,
|
||||
)
|
||||
continue
|
||||
|
||||
elif ctx.history:
|
||||
# --- FIRST REVERT (from OPTIMIZE state) ---
|
||||
ctx.local_retries = 0
|
||||
prev_cp = ctx.history.pop()
|
||||
ctx.revert_anchor_checkpoint = prev_cp
|
||||
_load_checkpoint(prev_cp)
|
||||
ctx.revert_anchor_step = int(np.round(mj_data.time / config.sim_dt))
|
||||
ctx.revert_attempts_from_anchor = 1
|
||||
# Reset boost so iters/noise don't compound across anchors —
|
||||
# each new anchor starts a fresh local-retry sweep at level 0.
|
||||
ctx.boost_level = 0
|
||||
_apply_boost(
|
||||
config,
|
||||
ctx.boost_level,
|
||||
orig_max_iters,
|
||||
orig_env_params_list,
|
||||
orig_noise_scale,
|
||||
)
|
||||
ctx.state = _StepState.RECOVERING
|
||||
loguru.logger.warning(
|
||||
"Revert attempt 1/{} from sim_step={} to anchor sim_step={}.",
|
||||
config.max_revert_depth,
|
||||
sim_step,
|
||||
ctx.revert_anchor_step,
|
||||
)
|
||||
continue
|
||||
|
||||
else:
|
||||
# --- DONE (no history to revert to) ---
|
||||
loguru.logger.error(
|
||||
"Unrecoverable failure at sim_step={} (no history).",
|
||||
sim_step,
|
||||
)
|
||||
_apply_boost(
|
||||
config, 0, orig_max_iters, orig_env_params_list, orig_noise_scale
|
||||
)
|
||||
_truncate_to_latest_commit()
|
||||
sim_step = _replay_and_advance(sim_step, infos)
|
||||
ctx.state = _StepState.DONE
|
||||
break
|
||||
|
||||
t1 = time.perf_counter()
|
||||
rtr = config.ctrl_dt / (t1 - t0)
|
||||
n_att = infos.get("num_attempts", [1])
|
||||
n_att_v = n_att[0] if hasattr(n_att, "__len__") else n_att
|
||||
print(
|
||||
f"Realtime rate: {rtr:.2f}, plan time: {t1 - t0:.4f}s, "
|
||||
f"sim_steps: {sim_step}/{config.max_sim_steps}, "
|
||||
f"opt_steps: {infos['opt_steps'][0]}, attempts: {n_att_v}, "
|
||||
f"state: {ctx.state}, boost: {ctx.boost_level}",
|
||||
end="\r",
|
||||
)
|
||||
|
||||
if sim_step >= config.max_sim_steps:
|
||||
ctx.state = _StepState.DONE
|
||||
break
|
||||
|
||||
t_end = time.perf_counter()
|
||||
print(f"Total time: {t_end - t_start:.4f}s")
|
||||
|
||||
if config.save_info and len(info_list) > 0:
|
||||
info_aggregated = {}
|
||||
for k in info_list[0]:
|
||||
if not isinstance(info_list[0][k], np.ndarray):
|
||||
continue
|
||||
vals = [info[k] for info in info_list if k in info]
|
||||
if len(vals) == 0:
|
||||
continue
|
||||
if all(v.shape == vals[0].shape for v in vals):
|
||||
info_aggregated[k] = np.stack(vals, axis=0)
|
||||
elif k in ("qpos", "qvel", "time", "ctrl"):
|
||||
info_aggregated[k] = np.concatenate(vals, axis=0)
|
||||
np.savez(
|
||||
f"{config.output_dir}/trajectory_mjwp_fast.npz",
|
||||
**info_aggregated,
|
||||
)
|
||||
loguru.logger.info(
|
||||
f"Saved info to {config.output_dir}/trajectory_mjwp_fast.npz"
|
||||
)
|
||||
|
||||
if config.save_video and len(images) > 0:
|
||||
video_path = f"{config.output_dir}/visualization_mjwp_fast.mp4"
|
||||
imageio.mimsave(
|
||||
video_path,
|
||||
images,
|
||||
fps=int(1 / config.render_dt),
|
||||
)
|
||||
loguru.logger.info(f"Saved video to {video_path}")
|
||||
|
||||
errors = None
|
||||
if info_list:
|
||||
qpos_traj = np.concatenate([info["qpos"] for info in info_list], axis=0)
|
||||
qpos_ref_np = qpos_ref[: qpos_traj.shape[0]].detach().cpu().numpy()
|
||||
errors = compute_object_tracking_error(
|
||||
qpos_traj, qpos_ref_np, config.embodiment_type, "mjwp"
|
||||
)
|
||||
loguru.logger.info(
|
||||
"Final object tracking error: pos={:.4f}, quat={:.4f}",
|
||||
errors["obj_pos_err"],
|
||||
errors["obj_quat_err"],
|
||||
)
|
||||
# Smoothness check: detect jumps between control ticks
|
||||
if qpos_traj.shape[0] > 1:
|
||||
diffs = np.linalg.norm(np.diff(qpos_traj, axis=0), axis=1)
|
||||
max_jump = diffs.max()
|
||||
max_jump_frame = int(diffs.argmax())
|
||||
loguru.logger.info(
|
||||
"Smoothness: max frame-to-frame qpos diff={:.6f} at frame {}/{}",
|
||||
max_jump,
|
||||
max_jump_frame,
|
||||
qpos_traj.shape[0],
|
||||
)
|
||||
for tick in range(config.ctrl_steps, qpos_traj.shape[0], config.ctrl_steps):
|
||||
boundary_jump = np.linalg.norm(qpos_traj[tick] - qpos_traj[tick - 1])
|
||||
if boundary_jump > 0.5:
|
||||
loguru.logger.warning(
|
||||
"Large jump at tick boundary frame {}: diff={:.4f}",
|
||||
tick,
|
||||
boundary_jump,
|
||||
)
|
||||
|
||||
if "viser" in config.viewer and config.save_viser:
|
||||
viser_path = f"{config.output_dir}/visualization_mjwp_fast.viser"
|
||||
try:
|
||||
from spider.viewers.viser_viewer import (
|
||||
_STATE as _viser_state,
|
||||
end_recording_and_save,
|
||||
)
|
||||
|
||||
# If a recording was started, prefer the timeline recording.
|
||||
# Falls back to a single-frame snapshot when no recording exists.
|
||||
ok = end_recording_and_save(viser_path)
|
||||
if ok:
|
||||
loguru.logger.info(
|
||||
f"Saved viser optimization recording to {viser_path}"
|
||||
)
|
||||
else:
|
||||
server = _viser_state.server
|
||||
if server is not None:
|
||||
Path(viser_path).write_bytes(
|
||||
server.get_scene_serializer().serialize()
|
||||
)
|
||||
loguru.logger.info(f"Saved viser scene to {viser_path}")
|
||||
except Exception as e:
|
||||
loguru.logger.warning(f"Failed to save viser scene: {e}")
|
||||
|
||||
if "viser" in config.viewer and config.wait_on_finish:
|
||||
loguru.logger.info(
|
||||
"Optimization complete! Keeping Viser server alive. Press Ctrl+C to exit."
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="config", config_name="default")
|
||||
def run_main(cfg: DictConfig) -> None:
|
||||
"""Entry point for Hydra configuration runner."""
|
||||
config_dict = dict(cfg)
|
||||
|
||||
load_config_path = config_dict.get("load_config_path", "")
|
||||
if load_config_path:
|
||||
loaded_config = load_config_yaml(load_config_path)
|
||||
cli_overrides = _extract_cli_overrides(cfg)
|
||||
config_dict = {**loaded_config, **cli_overrides}
|
||||
else:
|
||||
config_dict = filter_config_fields(config_dict)
|
||||
|
||||
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"])
|
||||
|
||||
config = Config(**config_dict)
|
||||
main(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_main()
|
||||
+766
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user