chomeed/mimicgen_square_d0_mtdit_flow_40k_robotstate9

mimicgen_square_d0 multi_task_dit (flow-matching, robotstate9) — step 40000

LeRobot multi_task_dit flow-matching policy for MimicGen Square_D0, trained for 40000 steps.

State: "robotstate9" via dataset slicing (9-D, object dropped)

Unlike chomeed/mimicgen_three_piece_assembly_d0_flow_matching_prior_40k_robotstate9 (which keeps the full 51-D observation.state and slices to the first 9 inside the policy via n_state_dims), this model was trained on a dataset whose observation.state was pre-sliced to the first 9 dims:

observation.state = [ eef_pos(3) | eef_quat(4) | gripper_qpos(2) ]   = 9

i.e. the 42/14 privileged object dims were dropped from the data. config.json declares observation.state with shape [9], so this loads on stock LeRobot multi_task_dit (no n_state_dims field needed). The policy is vision-only for object state (agentview + eye_in_hand).

Training

dataset	chomeed/mimicgen_square_d0_224x224, observation.state sliced to first 9 dims
policy	multi_task_dit, objective=flow_matching
steps	40000
batch size	64
seed	1000
optimizer	Adam, lr 2e-5, cosine schedule

Inference

n_obs_steps	2
n_action_steps	24
horizon	32
num_integration_steps	100 (euler)
cameras	agentview, eye_in_hand (each 3×224×224)
action dim	7 (OSC_POSE + gripper)

Eval

Evaluated on the Square_D0 robosuite env (headless OSMesa software rendering); success = env is_success reached at any step. Per-checkpoint success-rate curve: see square_d0_eval_success.png (added to this repo).

The CLIP text_model/vision_model ... UNEXPECTED warnings at load time are benign.