vla-jepa-folding / train_config.json
chalkp's picture
Upload folder using huggingface_hub
9449989 verified
Raw
History Blame Contribute Delete
8 kB
{
"dataset": {
"repo_id": "lerobot/high_quality_folding",
"root": "lerobot/high_quality_folding",
"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": true,
"video_backend": "torchcodec",
"return_uint8": false,
"streaming": false
},
"env": null,
"policy": {
"type": "vla_jepa",
"n_obs_steps": 1,
"input_features": {
"observation.state": {
"type": "STATE",
"shape": [
16
]
},
"observation.images.left_wrist": {
"type": "VISUAL",
"shape": [
3,
720,
1280
]
},
"observation.images.right_wrist": {
"type": "VISUAL",
"shape": [
3,
720,
1280
]
},
"observation.images.base": {
"type": "VISUAL",
"shape": [
3,
480,
640
]
}
},
"output_features": {
"action": {
"type": "ACTION",
"shape": [
16
]
}
},
"device": "cuda",
"use_amp": true,
"use_peft": false,
"push_to_hub": false,
"repo_id": "chalkp/vla-jepa-folding",
"private": false,
"tags": [
"vla-jepa",
"folding",
"bimanual",
"fine-tuned"
],
"license": "apache-2.0",
"pretrained_path": "chalkp/vla-jepa-folding",
"chunk_size": 7,
"n_action_steps": 7,
"normalization_mapping": {
"VISUAL": "IDENTITY",
"STATE": "QUANTILES",
"ACTION": "QUANTILES"
},
"qwen_model_name": "Qwen/Qwen3-VL-2B-Instruct",
"jepa_encoder_name": "facebook/vjepa2-vitl-fpc64-256",
"freeze_qwen": true,
"enable_world_model": false,
"reinit_modules": [
"model.action_model.action_encoder",
"model.action_model.action_decoder",
"model.action_model.state_encoder"
],
"tokenizer_padding_side": "left",
"prompt_template": "Your task is {instruction}. Infer the temporal dynamics from frames {actions} and produce the corresponding policy actions {e_actions}.",
"special_action_token": "<|action_{}|>",
"embodied_action_token": "<|embodied_action|>",
"action_dim": 16,
"state_dim": 16,
"num_action_tokens_per_timestep": 8,
"num_embodied_action_tokens_per_instruction": 32,
"num_inference_timesteps": 4,
"action_hidden_size": 1024,
"action_model_type": "DiT-B",
"action_num_layers": 16,
"action_num_heads": 12,
"action_attention_head_dim": 64,
"action_dropout": 0.2,
"action_num_timestep_buckets": 1000,
"action_noise_beta_alpha": 1.5,
"action_noise_beta_beta": 1.0,
"action_noise_s": 0.999,
"num_target_vision_tokens": 32,
"action_max_seq_len": 1024,
"num_video_frames": 8,
"predictor_depth": 12,
"predictor_num_heads": 8,
"predictor_mlp_ratio": 4.0,
"predictor_dropout": 0.0,
"world_model_loss_weight": 0.1,
"jepa_tubelet_size": 2,
"repeated_diffusion_steps": 8,
"resize_images_to": [
224,
224
],
"binarize_gripper_action": false,
"pre_snap_gripper_action": false,
"clip_normalized_actions": true,
"gripper_dim": 7,
"gripper_threshold": 0.5,
"torch_dtype": "bfloat16",
"optimizer_lr": 3.75e-05,
"optimizer_betas": [
0.9,
0.95
],
"optimizer_eps": 1e-08,
"optimizer_weight_decay": 0.01,
"optimizer_grad_clip_norm": 1.0,
"scheduler_warmup_steps": 1000,
"scheduler_decay_steps": 40000,
"scheduler_decay_lr": 2.5e-06
},
"reward_model": null,
"output_dir": "outputs/train/vla-jepa-folding",
"job_name": "vla-jepa-folding",
"resume": true,
"seed": 1000,
"cudnn_deterministic": false,
"num_workers": 4,
"batch_size": 8,
"prefetch_factor": 4,
"persistent_workers": true,
"steps": 40000,
"eval_freq": 10000,
"log_freq": 500,
"tolerance_s": 0.01,
"save_checkpoint": true,
"save_freq": 10000,
"use_policy_training_preset": true,
"optimizer": {
"type": "adamw",
"lr": 3.75e-05,
"weight_decay": 0.01,
"grad_clip_norm": 1.0,
"betas": [
0.9,
0.95
],
"eps": 1e-08
},
"scheduler": {
"type": "cosine_decay_with_warmup",
"num_warmup_steps": 1000,
"num_decay_steps": 40000,
"peak_lr": 3.75e-05,
"decay_lr": 2.5e-06
},
"eval": {
"n_episodes": 0,
"batch_size": 0,
"use_async_envs": false
},
"wandb": {
"enable": true,
"disable_artifact": true,
"project": "lerobot-folding",
"entity": null,
"notes": null,
"run_id": "ccjqh3dl",
"mode": "offline",
"add_tags": true
},
"peft": null,
"sample_weighting": {
"type": "rabc",
"progress_path": "lerobot/high_quality_folding/sarm_progress.parquet",
"head_mode": "sparse",
"kappa": 0.0265,
"epsilon": 1e-06,
"extra_params": {}
},
"rename_map": {},
"checkpoint_path": "outputs/train/vla-jepa-folding/checkpoints/last"
}