LeeHakHo/mimicgen_aux_id90_checkpoints

MimicGen ID90 — Diffusion Policy checkpoints, four auxiliary-head arms

Trained policies for a study of where an auxiliary object-pose head buys out-of-distribution generalization. Each task is trained on a rotation-restricted ID90 distribution and evaluated on OOD rungs that move the manipulated object's position and yaw, so a drop can be attributed to one factor rather than to "the scene changed".

Every checkpoint here is a Diffusion Policy (UNet, DDPM training / DDIM sampling at 10 inference steps) on 84x84 images from agentview and robot0_eye_in_hand, trained 1000 epochs on 200 generated demonstrations.

The four arms

All four share the backbone, the data and the schedule. They differ only in what the auxiliary head is asked to predict from the shared visual features — the policy loss is identical.

arm	auxiliary target
baseline	none (the head is absent)
aux_world_frame	the manipulated object's pose in the world frame
aux_eef_frame	the same pose expressed in the end-effector frame (position residual + relative yaw)
aux_obj_eef_frame	both, from one head: the absolute pose and the eef-frame pose concatenated

Rotation is supervised as (sin(N*yaw), cos(N*yaw)), with N the object's symmetry order.

Layout

<task>/<arm>/model_epoch_<N>.pth      N = 100, 200, ... 1000

Ten checkpoints per run, the same grid the evaluation uses. Training wrote one every 10 epochs; this repo carries every hundredth.

Tasks in this release

square_d2, threading_d0, three_piece_assembly_d0, mug_cleanup_d1, kitchen_d1, pick_place_d0 — 6 of the study's 12, x 4 arms = 24 runs, 240 checkpoints.

The other six are deliberately absent rather than merely unfinished, and are being regenerated:

• nut_assembly_d0 — its OOD ladder had no effect at all. The env never passed a placement_initializer, so robosuite built its own and the bounds dict the ID/OOD rules rewrite was read by nothing: ID, OOD_POS, OOD_YAW and OOD_BOTH sampled identically, and the training demos carried full-circle nut yaw instead of the 90-degree ID90 window.
• stack_d1, stack_three_d1 — their cubes are exact cubes, so C4 about z: the orientation orbit is 90 degrees and the 90-degree ID window already covered all of it, which made the yaw rung a resample of the training distribution rather than a shift.
• hammer_cleanup_d1 — its ID distribution had saturated (success 1.000), leaving no headroom to separate the arms.
• coffee_d2, coffee_preparation_d1 — retraining with the auxiliary rotation target masked per object. Their coffee pod is a body of revolution with a pinned yaw, so its rotation target was noise no image can predict.

Loading

Standard robomimic checkpoints — a dict with model, config, algo_name, env_metadata, shape_metadata, action_normalization_stats. Weights only; no optimizer state, so these are for evaluation and fine-tuning starts, not for resuming a run.

from robomimic.utils.file_utils import policy_from_checkpoint
policy, ckpt = policy_from_checkpoint(ckpt_path="square_d2/aux_obj_eef_frame/model_epoch_1000.pth")

Reproducing

The stack is pinned; two of the four packages are forks and the other two are upstream at a fixed commit.

package	source	commit
mimicgen	LeeHakHo/mimicgen_aux (fork)	main @ da00ece
robomimic	LeeHakHo/Robomimic_Async (fork)	feature/id90-12task @ cb24141
robosuite	ARISE-Initiative/robosuite	b9d8d3de
robosuite_task_zoo	ARISE-Initiative/robosuite-task-zoo	74eab7f

setup_env/ENVIRONMENT.md in the mimicgen fork carries the full recipe. Two things there are easy to miss: numpy must stay on 1.x (moving to 2.2.6 took a baseline from 0.29 to 0.02 at evaluation while training looked entirely normal), and robosuite needs two patched textures that pick_place_d0's 84x84 observations depend on.

Demonstrations and the fixed evaluation scenes: LeeHakHo/mimicgen_aux_data.

License

cc-by-nc-sa-4.0, inherited from the training data.