{ "dataset": { "repo_id": "shubham4413/so101_wm", "root": "/hpcwork/dl125352/hf/hub/datasets--shubham4413--so101_wm/snapshots/96e2ecc061a02d2952083128350f6815de30cb9f/", "episodes": [ 0, 1, 2, 3, 4, 6, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 29, 30, 31, 32, 33, 34, 35, 36, 37, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 50, 51, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 66, 67, 68, 69, 70, 71, 72, 75, 76, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 93, 94, 95, 96, 97, 98, 100, 101, 103, 104, 105, 106, 107, 108, 109, 111, 112, 113, 114, 116, 117, 118, 119, 120, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 139, 142, 143, 144, 145, 146, 147, 148, 149, 150, 152, 153, 154, 155, 156, 157, 158, 160, 161, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176 ], "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.exterior_1_left": { "type": "VISUAL", "shape": [ 3, 224, 224 ] }, "observation.images.exterior_2_left": { "type": "VISUAL", "shape": [ 3, 224, 224 ] } }, "output_features": { "action": { "type": "ACTION", "shape": [ 6 ] } }, "device": "cuda", "use_amp": false, "use_peft": true, "push_to_hub": false, "repo_id": "shubham4413/so101_wm_vlajepa", "private": null, "tags": null, "license": null, "pretrained_path": "lerobot/VLA-JEPA-Pretrain", "pretrained_revision": null, "chunk_size": 7, "n_action_steps": 7, "normalization_mapping": { "VISUAL": "IDENTITY", "STATE": "MEAN_STD", "ACTION": "MIN_MAX" }, "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": 6, "state_dim": 6, "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": true, "pre_snap_gripper_action": true, "clip_normalized_actions": true, "gripper_dim": 5, "gripper_threshold": 0.5, "torch_dtype": "bfloat16", "optimizer_lr": 0.0001, "optimizer_betas": [ 0.9, 0.95 ], "optimizer_eps": 1e-08, "optimizer_weight_decay": 1e-08, "optimizer_grad_clip_norm": 1.0, "scheduler_warmup_steps": 5000, "scheduler_decay_steps": 30000, "scheduler_decay_lr": 1e-06 }, "reward_model": null, "output_dir": "/hpcwork/dl125352/train/so101_wm_vlajepa", "job_name": "vla_jepa", "resume": false, "seed": 1000, "cudnn_deterministic": false, "num_workers": 8, "batch_size": 8, "prefetch_factor": 4, "persistent_workers": true, "steps": 30000, "env_eval_freq": 20000, "log_freq": 100, "eval_steps": 0, "max_eval_samples": 0, "tolerance_s": 0.001, "save_checkpoint": true, "save_freq": 10000, "use_policy_training_preset": true, "optimizer": { "type": "adamw", "lr": 0.0001, "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": 5000, "num_decay_steps": 30000, "peak_lr": 0.0001, "decay_lr": 1e-06 }, "eval": { "n_episodes": 50, "batch_size": 50, "use_async_envs": true, "recording": false, "recording_repo_id": null, "recording_private": false }, "wandb": { "enable": false, "disable_artifact": false, "project": "lerobot", "entity": null, "notes": null, "run_id": null, "mode": null, "add_tags": true }, "peft": { "target_modules": ".*model\\.qwen\\..*\\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)", "full_training_modules": [ "model.action_model", "model.video_predictor" ], "method_type": "LORA", "init_type": null, "r": 16, "lora_alpha": 32 }, "sample_weighting": null, "rename_map": { "observation.images.left": "observation.images.exterior_1_left", "observation.images.fpv": "observation.images.exterior_2_left" }, "checkpoint_path": null }