Robotics
LeRobot
Safetensors
fastwam
omx_multicubes_fastwam / train_config.json
maximellerbach's picture
Upload policy weights, train config and readme
807ab9b verified
Raw
History Blame Contribute Delete
7.62 kB
{
"dataset": {
"repo_id": "maximellerbach/omx_multicubes",
"root": null,
"episodes": null,
"image_transforms": {
"enable": false,
"max_num_transforms": 3,
"random_order": false,
"tfs": {
"brightness": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"brightness": [
0.8,
1.2
]
}
},
"contrast": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"contrast": [
0.8,
1.2
]
}
},
"saturation": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"saturation": [
0.5,
1.5
]
}
},
"hue": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"hue": [
-0.05,
0.05
]
}
},
"sharpness": {
"weight": 1.0,
"type": "SharpnessJitter",
"kwargs": {
"sharpness": [
0.5,
1.5
]
}
},
"affine": {
"weight": 1.0,
"type": "RandomAffine",
"kwargs": {
"degrees": [
-5.0,
5.0
],
"translate": [
0.05,
0.05
]
}
}
}
},
"revision": null,
"use_imagenet_stats": false,
"video_backend": "torchcodec",
"return_uint8": false,
"streaming": false
},
"env": null,
"policy": {
"type": "fastwam",
"n_obs_steps": 1,
"input_features": {
"observation.images.top": {
"type": "VISUAL",
"shape": [
3,
224,
224
]
},
"observation.images.wrist": {
"type": "VISUAL",
"shape": [
3,
224,
224
]
},
"observation.state": {
"type": "STATE",
"shape": [
6
]
}
},
"output_features": {
"action": {
"type": "ACTION",
"shape": [
6
]
}
},
"device": "cuda",
"use_amp": false,
"use_peft": false,
"push_to_hub": true,
"repo_id": "maximellerbach/omx_multicubes_fastwam",
"private": null,
"tags": null,
"license": null,
"pretrained_path": "lerobot/fastwam_base",
"pretrained_revision": null,
"action_dim": 6,
"proprio_dim": 6,
"action_horizon": 32,
"n_action_steps": 32,
"num_video_frames": 33,
"action_video_freq_ratio": 4,
"image_size": [
224,
448
],
"context_len": 128,
"model_id": "Wan-AI/Wan2.2-TI2V-5B",
"tokenizer_model_id": "Wan-AI/Wan2.2-TI2V-5B",
"base_model_id": "lerobot/fastwam_base",
"tokenizer_max_len": 128,
"load_text_encoder": true,
"mot_checkpoint_mixed_attn": false,
"torch_dtype": "bfloat16",
"prompt_template": "A video recorded from a robot's point of view executing the following instruction: {task}",
"num_inference_steps": 10,
"inference_seed": 42,
"rand_device": "cpu",
"text_cfg_scale": 1.0,
"negative_prompt": "",
"sigma_shift": null,
"tiled": false,
"fp32_attention": true,
"use_gradient_checkpointing": true,
"freeze_video_expert": false,
"toggle_action_dimensions": [],
"video_scheduler": {
"train_shift": 5.0,
"infer_shift": 5.0,
"num_train_timesteps": 1000
},
"action_scheduler": {
"train_shift": 5.0,
"infer_shift": 5.0,
"num_train_timesteps": 1000
},
"loss": {
"lambda_video": 1.0,
"lambda_action": 1.0
},
"video_dit_config": {
"patch_size": [
1,
2,
2
],
"in_dim": 48,
"hidden_dim": 3072,
"ffn_dim": 14336,
"freq_dim": 256,
"text_dim": 4096,
"out_dim": 48,
"num_heads": 24,
"attn_head_dim": 128,
"num_layers": 30,
"eps": 1e-06,
"seperated_timestep": true,
"use_gradient_checkpointing": true,
"video_attention_mask_mode": "first_frame_causal",
"action_conditioned": false,
"action_dim": 6,
"action_group_causal_mask_mode": "group_diagonal",
"fp32_attention": true
},
"action_dit_config": {
"action_dim": 6,
"hidden_dim": 1024,
"ffn_dim": 4096,
"num_heads": 24,
"attn_head_dim": 128,
"num_layers": 30,
"text_dim": 4096,
"freq_dim": 256,
"eps": 1e-06,
"use_gradient_checkpointing": true,
"fp32_attention": true
},
"normalization_mapping": {
"VISUAL": "IDENTITY",
"STATE": "MEAN_STD",
"ACTION": "MEAN_STD"
},
"optimizer_lr": 0.0001,
"optimizer_weight_decay": 0.01
},
"reward_model": null,
"output_dir": "outputs/train/omx_multicubes_fastwam_full",
"job_name": "fastwam",
"resume": false,
"seed": 1000,
"cudnn_deterministic": false,
"num_workers": 4,
"batch_size": 8,
"prefetch_factor": 4,
"persistent_workers": true,
"steps": 6000,
"eval_freq": 20000,
"log_freq": 50,
"tolerance_s": 0.0001,
"save_checkpoint": true,
"save_freq": 6000,
"use_policy_training_preset": true,
"optimizer": {
"type": "adamw",
"lr": 0.0001,
"weight_decay": 0.01,
"grad_clip_norm": 10.0,
"betas": [
0.9,
0.999
],
"eps": 1e-08
},
"scheduler": null,
"eval": {
"n_episodes": 50,
"batch_size": 50,
"use_async_envs": true
},
"wandb": {
"enable": true,
"disable_artifact": true,
"project": "lerobot",
"entity": null,
"notes": null,
"run_id": "9567pvf6",
"mode": null,
"add_tags": true
},
"peft": null,
"sample_weighting": null,
"rename_map": {},
"checkpoint_path": null
}