{ "dataset": { "repo_id": "lerobot-data-collection/level2_final_quality3_t_0_hil_data_c", "root": "/fsx/maxime_ellerbach/datasets/level2_final_quality3_t_0_hil_data_c_rel", "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": "pyav", "return_uint8": false, "depth_output_unit": "mm", "streaming": false, "eval_split": 0.0 }, "env": null, "policy": { "type": "vla_jepa", "n_obs_steps": 1, "input_features": { "observation.images.base": { "type": "VISUAL", "shape": [ 3, 224, 224 ] }, "observation.images.right_wrist": { "type": "VISUAL", "shape": [ 3, 224, 224 ] }, "observation.images.left_wrist": { "type": "VISUAL", "shape": [ 3, 224, 224 ] }, "observation.state": { "type": "STATE", "shape": [ 16 ] } }, "output_features": { "action": { "type": "ACTION", "shape": [ 16 ] } }, "device": "cuda", "use_amp": false, "use_peft": false, "push_to_hub": true, "repo_id": "maximellerbach/folding_vla_jepa", "private": null, "tags": null, "license": null, "pretrained_path": "/fsx/maxime_ellerbach/robots/outputs/train/folding_vla_jepa_multinode_3711/checkpoints/last/pretrained_model", "pretrained_revision": null, "chunk_size": 30, "n_action_steps": 30, "normalization_mapping": { "VISUAL": "IDENTITY", "STATE": "MEAN_STD", "ACTION": "MEAN_STD" }, "qwen_model_name": "Qwen/Qwen3-VL-2B-Instruct", "jepa_encoder_name": "facebook/vjepa2-vitl-fpc64-256", "freeze_qwen": false, "enable_world_model": true, "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, "use_relative_actions": true, "relative_exclude_joints": [ "gripper" ], "action_feature_names": [ "right_joint_1.pos", "right_joint_2.pos", "right_joint_3.pos", "right_joint_4.pos", "right_joint_5.pos", "right_joint_6.pos", "right_joint_7.pos", "right_gripper.pos", "left_joint_1.pos", "left_joint_2.pos", "left_joint_3.pos", "left_joint_4.pos", "left_joint_5.pos", "left_joint_6.pos", "left_joint_7.pos", "left_gripper.pos" ], "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": false, "gripper_dim": 6, "gripper_threshold": 0.5, "torch_dtype": "bfloat16", "optimizer_lr": 2e-05, "optimizer_betas": [ 0.9, 0.95 ], "optimizer_eps": 1e-08, "optimizer_weight_decay": 1e-08, "optimizer_grad_clip_norm": 1.0, "scheduler_warmup_steps": 500, "scheduler_decay_steps": 30000, "scheduler_decay_lr": 2e-05 }, "reward_model": null, "output_dir": "outputs/train/folding_vla_jepa_statefix_6400", "job_name": "folding_vla_jepa_statefix", "resume": false, "seed": 1000, "cudnn_deterministic": false, "num_workers": 4, "batch_size": 16, "prefetch_factor": 4, "persistent_workers": true, "steps": 5000, "env_eval_freq": 20000, "log_freq": 50, "eval_steps": 0, "max_eval_samples": 0, "tolerance_s": 0.0001, "save_checkpoint": true, "save_freq": 1000, "use_policy_training_preset": true, "optimizer": { "type": "adamw", "lr": 2e-05, "weight_decay": 1e-08, "grad_clip_norm": 1.0, "betas": [ 0.9, 0.95 ], "eps": 1e-08 }, "scheduler": { "type": "cosine_decay_with_warmup", "num_warmup_steps": 500, "num_decay_steps": 30000, "peak_lr": 2e-05, "decay_lr": 2e-05 }, "eval": { "n_episodes": 50, "batch_size": 50, "use_async_envs": true, "recording": false, "recording_repo_id": null, "recording_private": false, "save_predicted_video": false }, "wandb": { "enable": true, "disable_artifact": true, "project": "folding-tshirt", "entity": null, "notes": null, "run_id": "i7jnld97", "mode": null, "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": { "type": "rabc", "progress_path": "hf://datasets/lerobot-data-collection/level2_final_quality3_t_0_hil_data_c/sarm_progress.parquet", "head_mode": "sparse", "kappa": 0.05, "epsilon": 1e-06, "extra_params": {} }, "rename_map": {}, "checkpoint_path": null }