robotwin-hanging-mug-dit / train_config.json
Neotix-Robotics's picture
multi-task DiT trained 50k steps on RoboTwin
d49ccf6 verified
Raw History Blame Contribute Delete
8.4 kB
{
"dataset": {
"repo_id": "lerobot/robotwin_unified",
"repo_type": "dataset",
"root": null,
"episodes": null,
"exclude_episodes": null,
"image_transforms": {
"enable": true,
"max_num_transforms": 4,
"random_order": false,
"tfs": {
"brightness": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"brightness": [
0.75,
1.25
]
}
},
"contrast": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"contrast": [
0.6,
1.4
]
}
},
"saturation": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"saturation": [
0.8,
1.2
]
}
},
"hue": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"hue": [
-0.05,
0.05
]
}
}
}
},
"revision": "main",
"use_imagenet_stats": true,
"video_backend": "torchcodec",
"return_uint8": false,
"depth_output_unit": "mm",
"streaming": false,
"eval_split": 0.0
},
"env": null,
"policy": {
"type": "multi_task_dit",
"n_obs_steps": 2,
"input_features": {
"observation.state": {
"type": "STATE",
"shape": [
14
]
},
"observation.images.cam_high": {
"type": "VISUAL",
"shape": [
480,
640,
3
]
},
"observation.images.cam_left_wrist": {
"type": "VISUAL",
"shape": [
480,
640,
3
]
},
"observation.images.cam_right_wrist": {
"type": "VISUAL",
"shape": [
480,
640,
3
]
}
},
"output_features": {
"action": {
"type": "ACTION",
"shape": [
14
]
}
},
"device": "cuda",
"use_amp": true,
"use_peft": false,
"push_to_hub": true,
"repo_id": "unitzero-ai/dit_robotwin_2026-09-29_15-41-09",
"private": true,
"tags": null,
"license": null,
"pretrained_path": "outputs/train/2026-09-30/09-34-48_dit_robotwin_2026-09-29_15-41-09/checkpoints/012000/pretrained_model",
"pretrained_revision": null,
"horizon": 32,
"n_action_steps": 24,
"objective": "diffusion",
"noise_scheduler_type": "DDPM",
"num_train_timesteps": 100,
"beta_schedule": "squaredcos_cap_v2",
"beta_start": 0.0001,
"beta_end": 0.02,
"prediction_type": "epsilon",
"clip_sample": true,
"clip_sample_range": 1.0,
"num_inference_steps": 10,
"sigma_min": 0.0,
"num_integration_steps": 100,
"integration_method": "euler",
"timestep_sampling_strategy": "beta",
"timestep_sampling_s": 0.999,
"timestep_sampling_alpha": 1.5,
"timestep_sampling_beta": 1.0,
"hidden_dim": 768,
"num_layers": 8,
"num_heads": 12,
"dropout": 0.1,
"use_positional_encoding": false,
"timestep_embed_dim": 256,
"use_rope": true,
"rope_base": 10000.0,
"vision_encoder_name": "openai/clip-vit-base-patch16",
"use_separate_rgb_encoder_per_camera": false,
"vision_encoder_lr_multiplier": 0.1,
"image_resize_shape": [
256,
256
],
"image_crop_shape": [
224,
224
],
"image_crop_is_random": true,
"text_encoder_name": "openai/clip-vit-base-patch16",
"tokenizer_max_length": 77,
"tokenizer_padding": "max_length",
"tokenizer_padding_side": "right",
"tokenizer_truncation": true,
"normalization_mapping": {
"VISUAL": "MEAN_STD",
"STATE": "MIN_MAX",
"ACTION": "MIN_MAX"
},
"optimizer_lr": 0.0003,
"optimizer_betas": [
0.95,
0.999
],
"optimizer_eps": 1e-08,
"optimizer_weight_decay": 0.0,
"scheduler_name": "cosine",
"scheduler_warmup_steps": 0,
"do_mask_loss_for_padding": false,
"drop_n_last_frames": 7
},
"reward_model": null,
"output_dir": "outputs/train/2026-09-30/09-34-48_dit_robotwin_2026-09-29_15-41-09",
"job_name": "dit_robotwin",
"resume": true,
"seed": 123,
"cudnn_deterministic": false,
"num_workers": 8,
"batch_size": 64,
"prefetch_factor": 4,
"persistent_workers": true,
"dataloader_multiprocessing_context": "spawn",
"steps": 50000,
"env_eval_freq": 0,
"log_freq": 100,
"eval_steps": 0,
"max_eval_samples": 0,
"tolerance_s": 0.0001,
"save_checkpoint": true,
"save_freq": 4000,
"checkpoint_format": "safetensors",
"use_policy_training_preset": true,
"optimizer": {
"type": "adam",
"lr": 0.0003,
"weight_decay": 0.0,
"grad_clip_norm": 10.0,
"betas": [
0.95,
0.999
],
"eps": 1e-08
},
"scheduler": {
"type": "diffuser",
"num_warmup_steps": 0,
"name": "cosine"
},
"parallelism": {
"dp_replicate": 8,
"dp_shard": 1,
"context_parallel": {
"ring_degree": 1,
"ulysses_degree": 1
},
"cfg_parallel": 1
},
"accelerator": {
"mixed_precision": "bf16",
"gradient_accumulation": {
"steps": 1
},
"fsdp": {
"reshard_after_forward": true,
"wrap_modules": null,
"min_num_params": null,
"cpu_offload": false,
"ignored_modules": null
},
"ddp": {
"find_unused_parameters": true,
"gradient_as_bucket_view": false,
"static_graph": false
},
"compile": {
"enabled": false,
"backend": "inductor",
"mode": null,
"regional": true
},
"activation_checkpointing": {
"mode": "none"
}
},
"eval": {
"n_episodes": 50,
"batch_size": 50,
"use_async_envs": true,
"recording": false,
"recording_repo_id": null,
"recording_private": false
},
"ema": {
"enable": false,
"decay": null,
"update_after_step": 0,
"inv_gamma": 1.0,
"power": 0.75,
"min_decay": 0.0,
"max_decay": 0.9999,
"use_for_eval": true
},
"wandb": {
"enable": true,
"disable_artifact": false,
"project": "dit_robotwin",
"entity": "unitzero-ai",
"notes": null,
"run_id": "qhy4lauf",
"resume": null,
"mode": null,
"console": "wrap",
"console_multipart": false,
"console_chunk_max_seconds": 0,
"add_tags": true
},
"peft": null,
"job": {
"target": null,
"image": "huggingface/lerobot-gpu:latest",
"timeout": "2d",
"detach": false,
"tags": []
},
"save_checkpoint_to_hub": false,
"sample_weighting": null,
"rename_map": {},
"checkpoint_path": "outputs/train/2026-09-30/09-34-48_dit_robotwin_2026-09-29_15-41-09/checkpoints/012000"
}