Upload policy weights, train config and readme

README.md

@@ -79,10 +79,10 @@ The policy consumes these observation features and produces these action feature
 
 | Setting | Value |
 | --- | --- |
-| Training steps | 5000 |
+| Training steps | 66686 |
 | Batch size | 16 |
 | Optimizer | adamw |
-| Learning rate | 2e-05 |
+| Learning rate | 0.0001 |
 | Seed | 1000 |
 | LeRobot version | 0.6.1 |
 

config.json

@@ -49,7 +49,7 @@
     "private": null,
     "tags": null,
     "license": null,
-    "pretrained_path": "/fsx/maxime_ellerbach/robots/outputs/train/folding_vla_jepa_multinode_3711/checkpoints/last/pretrained_model",
+    "pretrained_path": "/fsx/maxime_ellerbach/robots/runs/folding-tshirt/vla_jepa_base_3cam",
     "pretrained_revision": null,
     "chunk_size": 30,
     "n_action_steps": 30,
@@ -108,8 +108,8 @@
     "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_target_vision_tokens": 32,
     "num_video_frames": 8,
     "predictor_depth": 12,
     "predictor_num_heads": 8,
@@ -117,6 +117,7 @@
     "predictor_dropout": 0.0,
     "world_model_loss_weight": 0.1,
     "jepa_tubelet_size": 2,
+    "world_model_num_views": null,
     "repeated_diffusion_steps": 8,
     "resize_images_to": [
         224,
@@ -127,8 +128,11 @@
     "clip_normalized_actions": false,
     "gripper_dim": 6,
     "gripper_threshold": 0.5,
+    "gripper_joint_names": [
+        "gripper"
+    ],
     "torch_dtype": "bfloat16",
-    "optimizer_lr": 2e-05,
+    "optimizer_lr": 0.0001,
     "optimizer_betas": [
         0.9,
         0.95
@@ -136,7 +140,7 @@
     "optimizer_eps": 1e-08,
     "optimizer_weight_decay": 1e-08,
     "optimizer_grad_clip_norm": 1.0,
-    "scheduler_warmup_steps": 500,
+    "scheduler_warmup_steps": 5000,
     "scheduler_decay_steps": 30000,
     "scheduler_decay_lr": 2e-05
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:91835efb3763a07288ff882fe96ae7922f8a0ea4f42a84571a1f5ae8472e8752
-size 6174639408
+oid sha256:b2ab63be325c5f270d55508b74e346dfb90f108a89c97cfb363caf340e3313bf
+size 6157857616

train_config.json

@@ -134,7 +134,7 @@
         "private": null,
         "tags": null,
         "license": null,
-        "pretrained_path": "/fsx/maxime_ellerbach/robots/outputs/train/folding_vla_jepa_multinode_3711/checkpoints/last/pretrained_model",
+        "pretrained_path": "/fsx/maxime_ellerbach/robots/runs/folding-tshirt/vla_jepa_base_3cam",
         "pretrained_revision": null,
         "chunk_size": 30,
         "n_action_steps": 30,
@@ -193,8 +193,8 @@
         "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_target_vision_tokens": 32,
         "num_video_frames": 8,
         "predictor_depth": 12,
         "predictor_num_heads": 8,
@@ -202,6 +202,7 @@
         "predictor_dropout": 0.0,
         "world_model_loss_weight": 0.1,
         "jepa_tubelet_size": 2,
+        "world_model_num_views": null,
         "repeated_diffusion_steps": 8,
         "resize_images_to": [
             224,
@@ -212,8 +213,11 @@
         "clip_normalized_actions": false,
         "gripper_dim": 6,
         "gripper_threshold": 0.5,
+        "gripper_joint_names": [
+            "gripper"
+        ],
         "torch_dtype": "bfloat16",
-        "optimizer_lr": 2e-05,
+        "optimizer_lr": 0.0001,
         "optimizer_betas": [
             0.9,
             0.95
@@ -221,32 +225,34 @@
         "optimizer_eps": 1e-08,
         "optimizer_weight_decay": 1e-08,
         "optimizer_grad_clip_norm": 1.0,
-        "scheduler_warmup_steps": 500,
+        "scheduler_warmup_steps": 5000,
         "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",
+    "output_dir": "outputs/train/folding_vla_jepa_refactor_lr_8799",
+    "job_name": "folding_vla_jepa_refactor_lr",
     "resume": false,
+    "rebuild_processors": false,
     "seed": 1000,
     "cudnn_deterministic": false,
     "num_workers": 4,
     "batch_size": 16,
     "prefetch_factor": 4,
     "persistent_workers": true,
-    "steps": 5000,
+    "dataloader_multiprocessing_context": "spawn",
+    "steps": 66686,
     "env_eval_freq": 20000,
     "log_freq": 50,
     "eval_steps": 0,
     "max_eval_samples": 0,
     "tolerance_s": 0.0001,
     "save_checkpoint": true,
-    "save_freq": 1000,
+    "save_freq": 10000,
     "use_policy_training_preset": true,
     "optimizer": {
         "type": "adamw",
-        "lr": 2e-05,
+        "lr": 0.0001,
         "weight_decay": 1e-08,
         "grad_clip_norm": 1.0,
         "betas": [
@@ -257,9 +263,9 @@
     },
     "scheduler": {
         "type": "cosine_decay_with_warmup",
-        "num_warmup_steps": 500,
+        "num_warmup_steps": 5000,
         "num_decay_steps": 30000,
-        "peak_lr": 2e-05,
+        "peak_lr": 0.0001,
         "decay_lr": 2e-05
     },
     "eval": {
@@ -277,7 +283,7 @@
         "project": "folding-tshirt",
         "entity": null,
         "notes": null,
-        "run_id": "i7jnld97",
+        "run_id": "gsj0i9xi",
         "mode": null,
         "add_tags": true
     },