Upload folder using huggingface_hub

config.json

@@ -1,12 +1,28 @@
 {
-  "chunk_size": 100,
+  "chunk_size": 20,
   "dim_feedforward": 3200,
-  "dim_model": 512,
+  "dim_model": 64,
   "dropout": 0.1,
   "feedforward_activation": "relu",
+  "input_features": {
+    "observation.images.top": {
+      "shape": [
+        3,
+        480,
+        640
+      ],
+      "type": "VISUAL"
+    },
+    "observation.state": {
+      "shape": [
+        14
+      ],
+      "type": "STATE"
+    }
+  },
   "kl_weight": 10.0,
   "latent_dim": 32,
-  "n_action_steps": 100,
+  "n_action_steps": 20,
   "n_decoder_layers": 1,
   "n_encoder_layers": 4,
   "n_heads": 8,
@@ -20,6 +36,14 @@
   "optimizer_lr": 1e-05,
   "optimizer_lr_backbone": 1e-05,
   "optimizer_weight_decay": 0.0001,
+  "output_features": {
+    "action": {
+      "shape": [
+        14
+      ],
+      "type": "ACTION"
+    }
+  },
   "pre_norm": false,
   "pretrained_backbone_weights": "ResNet18_Weights.IMAGENET1K_V1",
   "replace_final_stride_with_dilation": false,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:81db69111d49e02f0adec79c9217d76c9b3e1d3df11574b76c33790aea07b18c
-size 206766560
+oid sha256:39628cc4c2accfc30bed32ed52c244e2a608f6d10d2122d69cf0f4f328e431b0
+size 60492680

train_config.json

@@ -1,9 +1,11 @@
 {
     "dataset": {
         "repo_id": "lerobot/aloha_sim_transfer_cube_human",
-        "episodes": null,
+        "episodes": [
+            0
+        ],
         "image_transforms": {
-            "enable": false,
+            "enable": true,
             "max_num_transforms": 3,
             "random_order": false,
             "tfs": {
@@ -66,16 +68,37 @@
     "env": {
         "type": "aloha",
         "n_envs": null,
-        "task": "AlohaTransferCube-v0",
+        "task": "AlohaInsertion-v0",
         "fps": 50,
-        "feature_types": {
-            "agent_pos": "STATE",
-            "pixels": {
-                "top": "VISUAL"
+        "features": {
+            "action": {
+                "type": "ACTION",
+                "shape": [
+                    14
+                ]
+            },
+            "agent_pos": {
+                "type": "STATE",
+                "shape": [
+                    14
+                ]
             },
-            "action": "ACTION"
+            "pixels/top": {
+                "type": "VISUAL",
+                "shape": [
+                    480,
+                    640,
+                    3
+                ]
+            }
         },
-        "episode_length": 400,
+        "features_map": {
+            "action": "action",
+            "agent_pos": "observation.state",
+            "top": "observation.image.top",
+            "pixels/top": "observation.images.top"
+        },
+        "episode_length": 5,
         "obs_type": "pixels_agent_pos",
         "render_mode": "rgb_array"
     },
@@ -87,13 +110,37 @@
             "STATE": "MEAN_STD",
             "ACTION": "MEAN_STD"
         },
-        "chunk_size": 100,
-        "n_action_steps": 100,
+        "input_features": {
+            "observation.images.top": {
+                "type": "VISUAL",
+                "shape": [
+                    3,
+                    480,
+                    640
+                ]
+            },
+            "observation.state": {
+                "type": "STATE",
+                "shape": [
+                    14
+                ]
+            }
+        },
+        "output_features": {
+            "action": {
+                "type": "ACTION",
+                "shape": [
+                    14
+                ]
+            }
+        },
+        "chunk_size": 20,
+        "n_action_steps": 20,
         "vision_backbone": "resnet18",
         "pretrained_backbone_weights": "ResNet18_Weights.IMAGENET1K_V1",
         "replace_final_stride_with_dilation": false,
         "pre_norm": false,
-        "dim_model": 512,
+        "dim_model": 64,
         "n_heads": 8,
         "dim_feedforward": 3200,
         "feedforward_activation": "relu",
@@ -112,17 +159,17 @@
     "output_dir": "outputs/train/act_aloha_sim_transfer_cube_human",
     "job_name": "aloha_act",
     "resume": false,
-    "device": "mps",
+    "device": "cpu",
     "use_amp": false,
     "seed": 1000,
     "num_workers": 4,
-    "batch_size": 8,
-    "eval_freq": 10,
+    "batch_size": 2,
+    "eval_freq": 20000,
     "log_freq": 1,
     "save_checkpoint": true,
-    "save_freq": 5,
+    "save_freq": 2,
     "offline": {
-        "steps": 10
+        "steps": 4
     },
     "online": {
         "steps": 0,
@@ -147,13 +194,9 @@
         "weight_decay": 0.0001,
         "grad_clip_norm": 10.0
     },
-    "scheduler": {
-        "type": "none",
-        "num_warmup_steps": 0
-    },
+    "scheduler": null,
     "eval": {
         "n_episodes": 1,
-        "episode_length": null,
         "batch_size": 1,
         "use_async_envs": false
     },