remove_hydra

README.md

@@ -1,7 +1,12 @@
 ---
 license: apache-2.0
 datasets:
-- JayLee131/vqbet_pusht
+- lerobot/pusht
+tags:
+- vqbet-policy
+- model_hub_mixin
+- pytorch_model_hub_mixin
+- robotics
 pipeline_tag: robotics
 ---
 # Model Card for VQ-BeT/PushT
@@ -14,21 +19,27 @@ See the [LeRobot library](https://github.com/huggingface/lerobot) (particularly
 
 ## Training Details
 
-Trained with [LeRobot@342f429](https://github.com/huggingface/lerobot/tree/342f429f1c321a2b4501c3007b1dacba7244b469).
+Trained with [LeRobot@3c0a209](https://github.com/huggingface/lerobot/tree/3c0a209f9fac4d2a57617e686a7f2a2309144ba2).
 
-The model was trained using this command:
+The model was trained using [LeRobot's training script](https://github.com/huggingface/lerobot/blob/main/lerobot/scripts/train.py) and with the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using this command:
 
 ```bash
 python lerobot/scripts/train.py \
-  policy=vqbet \
-  env=pusht dataset_repo_id=lerobot/pusht \
-  wandb.enable=true \
-  device=cuda
+    --output_dir=outputs/train/vqbet_pusht \
+    --policy.type=vqbet \
+    --dataset.repo_id=lerobot/pusht \
+    --env.type=pusht \
+    --seed=100000 \
+    --batch_size=64 \
+    --offline.steps=250000 \
+    --eval_freq=25000 \
+    --save_freq=25000 \
+    --wandb.enable=true
 ```
 
-The training curves may be found at https://wandb.ai/jaylee0301/lerobot/runs/9r0ndphr?nw=nwuserjaylee0301.
+The training curves may be found at https://wandb.ai/aliberts/lerobot/runs/3i7zs94u.
+The current model corresponds to the checkpoint at 200k steps.
 
-Training VQ-BeT on PushT took about 7-8 hours to train on an Nvida A6000.
 
 ## Model Size
 
@@ -51,4 +62,15 @@ Metric|Value
 Average max. overlap ratio for 500 episodes | 0.895
 Success rate for 500 episodes (%) | 63.8
 
-The results of each of the individual rollouts may be found in [eval_info.json](eval_info.json).
+The results of each of the individual rollouts may be found in [eval_info.json](eval_info.json). It was produced after training with this command:
+```bash
+python lerobot/scripts/eval.py \
+    --policy.path=outputs/train/vqbet_pusht/checkpoints/200000/pretrained_model \
+    --output_dir=outputs/eval/vqbet_pusht/200000 \
+    --env.type=pusht \
+    --seed=100000 \
+    --eval.n_episodes=500 \
+    --eval.batch_size=50 \
+    --device=cuda \
+    --use_amp=false
+```

config.json

@@ -1,53 +1,71 @@
 {
-  "action_chunk_size": 5,
-  "bet_softmax_temperature": 0.01,
-  "crop_is_random": true,
-  "crop_shape": [
-    84,
-    84
-  ],
-  "dropout": 0.1,
-  "gpt_block_size": 500,
-  "gpt_hidden_dim": 512,
-  "gpt_input_dim": 512,
-  "gpt_n_head": 8,
-  "gpt_n_layer": 8,
-  "gpt_output_dim": 512,
-  "input_normalization_modes": {
-    "observation.image": "mean_std",
-    "observation.state": "min_max"
-  },
-  "input_shapes": {
-    "observation.image": [
-      3,
-      96,
-      96
+    "type": "vqbet",
+    "n_obs_steps": 5,
+    "normalization_mapping": {
+        "VISUAL": "IDENTITY",
+        "STATE": "MIN_MAX",
+        "ACTION": "MIN_MAX"
+    },
+    "input_features": {
+        "observation.image": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                96,
+                96
+            ]
+        },
+        "observation.state": {
+            "type": "STATE",
+            "shape": [
+                2
+            ]
+        }
+    },
+    "output_features": {
+        "action": {
+            "type": "ACTION",
+            "shape": [
+                2
+            ]
+        }
+    },
+    "n_action_pred_token": 3,
+    "action_chunk_size": 5,
+    "vision_backbone": "resnet18",
+    "crop_shape": [
+        84,
+        84
     ],
-    "observation.state": [
-      2
-    ]
-  },
-  "mlp_hidden_dim": 1024,
-  "n_action_pred_token": 7,
-  "n_obs_steps": 5,
-  "n_vqvae_training_steps": 20000,
-  "offset_loss_weight": 10000.0,
-  "output_normalization_modes": {
-    "action": "min_max"
-  },
-  "output_shapes": {
-    "action": [
-      2
-    ]
-  },
-  "pretrained_backbone_weights": null,
-  "primary_code_loss_weight": 5.0,
-  "secondary_code_loss_weight": 0.5,
-  "sequentially_select": false,
-  "spatial_softmax_num_keypoints": 32,
-  "use_group_norm": true,
-  "vision_backbone": "resnet18",
-  "vqvae_embedding_dim": 256,
-  "vqvae_enc_hidden_dim": 128,
-  "vqvae_n_embed": 16
+    "crop_is_random": true,
+    "pretrained_backbone_weights": null,
+    "use_group_norm": true,
+    "spatial_softmax_num_keypoints": 32,
+    "n_vqvae_training_steps": 20000,
+    "vqvae_n_embed": 16,
+    "vqvae_embedding_dim": 256,
+    "vqvae_enc_hidden_dim": 128,
+    "gpt_block_size": 500,
+    "gpt_input_dim": 512,
+    "gpt_output_dim": 512,
+    "gpt_n_layer": 8,
+    "gpt_n_head": 8,
+    "gpt_hidden_dim": 512,
+    "dropout": 0.1,
+    "mlp_hidden_dim": 1024,
+    "offset_loss_weight": 10000.0,
+    "primary_code_loss_weight": 5.0,
+    "secondary_code_loss_weight": 0.5,
+    "bet_softmax_temperature": 0.1,
+    "sequentially_select": false,
+    "optimizer_lr": 0.0001,
+    "optimizer_betas": [
+        0.95,
+        0.999
+    ],
+    "optimizer_eps": 1e-08,
+    "optimizer_weight_decay": 1e-06,
+    "optimizer_vqvae_lr": 0.001,
+    "optimizer_vqvae_weight_decay": 0.0001,
+    "scheduler_warmup_steps": 500
 }

config.yaml

@@ -1,184 +0,0 @@
-resume: false
-device: cuda
-use_amp: false
-seed: 100000
-dataset_repo_id: lerobot/pusht
-video_backend: pyav
-training:
-  offline_steps: 250000
-  online_steps: 0
-  online_steps_between_rollouts: 1
-  online_sampling_ratio: 0.5
-  online_env_seed: ???
-  eval_freq: 20000
-  log_freq: 250
-  save_checkpoint: true
-  save_freq: 20000
-  num_workers: 4
-  batch_size: 64
-  image_transforms:
-    enable: false
-    max_num_transforms: 3
-    random_order: false
-    brightness:
-      weight: 1
-      min_max:
-      - 0.8
-      - 1.2
-    contrast:
-      weight: 1
-      min_max:
-      - 0.8
-      - 1.2
-    saturation:
-      weight: 1
-      min_max:
-      - 0.5
-      - 1.5
-    hue:
-      weight: 1
-      min_max:
-      - -0.05
-      - 0.05
-    sharpness:
-      weight: 1
-      min_max:
-      - 0.8
-      - 1.2
-  grad_clip_norm: 10
-  lr: 0.0001
-  lr_scheduler: cosine
-  lr_warmup_steps: 500
-  adam_betas:
-  - 0.95
-  - 0.999
-  adam_eps: 1.0e-08
-  adam_weight_decay: 1.0e-06
-  vqvae_lr: 0.001
-  n_vqvae_training_steps: 20000
-  bet_weight_decay: 0.0002
-  bet_learning_rate: 5.5e-05
-  bet_betas:
-  - 0.9
-  - 0.999
-  delta_timestamps:
-    observation.image:
-    - -0.4
-    - -0.3
-    - -0.2
-    - -0.1
-    - 0.0
-    observation.state:
-    - -0.4
-    - -0.3
-    - -0.2
-    - -0.1
-    - 0.0
-    action:
-    - -0.4
-    - -0.3
-    - -0.2
-    - -0.1
-    - 0.0
-    - 0.1
-    - 0.2
-    - 0.3
-    - 0.4
-    - 0.5
-    - 0.6
-    - 0.7
-    - 0.8
-    - 0.9
-    - 1.0
-eval:
-  n_episodes: 50
-  batch_size: 50
-  use_async_envs: false
-wandb:
-  enable: true
-  disable_artifact: false
-  project: lerobot
-  notes: ''
-fps: 10
-env:
-  name: pusht
-  task: PushT-v0
-  image_size: 96
-  state_dim: 2
-  action_dim: 2
-  fps: ${fps}
-  episode_length: 300
-  gym:
-    obs_type: pixels_agent_pos
-    render_mode: rgb_array
-    visualization_width: 384
-    visualization_height: 384
-override_dataset_stats:
-  observation.image:
-    mean:
-    - - - 0.5
-    - - - 0.5
-    - - - 0.5
-    std:
-    - - - 0.5
-    - - - 0.5
-    - - - 0.5
-  observation.state:
-    min:
-    - 13.456424
-    - 32.938293
-    max:
-    - 496.14618
-    - 510.9579
-  action:
-    min:
-    - 12.0
-    - 25.0
-    max:
-    - 511.0
-    - 511.0
-policy:
-  name: vqbet
-  n_obs_steps: 5
-  n_action_pred_token: 7
-  action_chunk_size: 5
-  input_shapes:
-    observation.image:
-    - 3
-    - 96
-    - 96
-    observation.state:
-    - ${env.state_dim}
-  output_shapes:
-    action:
-    - ${env.action_dim}
-  input_normalization_modes:
-    observation.image: mean_std
-    observation.state: min_max
-  output_normalization_modes:
-    action: min_max
-  vision_backbone: resnet18
-  crop_shape:
-  - 84
-  - 84
-  crop_is_random: true
-  pretrained_backbone_weights: null
-  use_group_norm: true
-  spatial_softmax_num_keypoints: 32
-  n_vqvae_training_steps: ${training.n_vqvae_training_steps}
-  vqvae_n_embed: 16
-  vqvae_embedding_dim: 256
-  vqvae_enc_hidden_dim: 128
-  gpt_block_size: 500
-  gpt_input_dim: 512
-  gpt_output_dim: 512
-  gpt_n_layer: 8
-  gpt_n_head: 8
-  gpt_hidden_dim: 512
-  dropout: 0.1
-  mlp_hidden_dim: 1024
-  offset_loss_weight: 10000.0
-  primary_code_loss_weight: 5.0
-  secondary_code_loss_weight: 0.5
-  bet_softmax_temperature: 0.01
-  sequentially_select: false

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:85995b992d83133f65e2a3d81bbeefeaa9e082f4e79d6cff5d5c5ab05cfc096b
-size 158152442
+oid sha256:f60b22049b275c026159fd4fcc018721ffbec47119ecfab98dc84814eb4e0d30
+size 158152042

train_config.json

@@ -0,0 +1,231 @@
+{
+    "dataset": {
+        "repo_id": "lerobot/pusht",
+        "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
+                        ]
+                    }
+                }
+            }
+        },
+        "local_files_only": false,
+        "use_imagenet_stats": true,
+        "video_backend": "pyav"
+    },
+    "env": {
+        "type": "pusht",
+        "task": "PushT-v0",
+        "fps": 10,
+        "features": {
+            "action": {
+                "type": "ACTION",
+                "shape": [
+                    2
+                ]
+            },
+            "agent_pos": {
+                "type": "STATE",
+                "shape": [
+                    2
+                ]
+            },
+            "pixels": {
+                "type": "VISUAL",
+                "shape": [
+                    384,
+                    384,
+                    3
+                ]
+            }
+        },
+        "features_map": {
+            "action": "action",
+            "agent_pos": "observation.state",
+            "environment_state": "observation.environment_state",
+            "pixels": "observation.image"
+        },
+        "episode_length": 300,
+        "obs_type": "pixels_agent_pos",
+        "render_mode": "rgb_array",
+        "visualization_width": 384,
+        "visualization_height": 384
+    },
+    "policy": {
+        "type": "vqbet",
+        "n_obs_steps": 5,
+        "normalization_mapping": {
+            "VISUAL": "IDENTITY",
+            "STATE": "MIN_MAX",
+            "ACTION": "MIN_MAX"
+        },
+        "input_features": {
+            "observation.image": {
+                "type": "VISUAL",
+                "shape": [
+                    3,
+                    96,
+                    96
+                ]
+            },
+            "observation.state": {
+                "type": "STATE",
+                "shape": [
+                    2
+                ]
+            }
+        },
+        "output_features": {
+            "action": {
+                "type": "ACTION",
+                "shape": [
+                    2
+                ]
+            }
+        },
+        "n_action_pred_token": 3,
+        "action_chunk_size": 5,
+        "vision_backbone": "resnet18",
+        "crop_shape": [
+            84,
+            84
+        ],
+        "crop_is_random": true,
+        "pretrained_backbone_weights": null,
+        "use_group_norm": true,
+        "spatial_softmax_num_keypoints": 32,
+        "n_vqvae_training_steps": 20000,
+        "vqvae_n_embed": 16,
+        "vqvae_embedding_dim": 256,
+        "vqvae_enc_hidden_dim": 128,
+        "gpt_block_size": 500,
+        "gpt_input_dim": 512,
+        "gpt_output_dim": 512,
+        "gpt_n_layer": 8,
+        "gpt_n_head": 8,
+        "gpt_hidden_dim": 512,
+        "dropout": 0.1,
+        "mlp_hidden_dim": 1024,
+        "offset_loss_weight": 10000.0,
+        "primary_code_loss_weight": 5.0,
+        "secondary_code_loss_weight": 0.5,
+        "bet_softmax_temperature": 0.1,
+        "sequentially_select": false,
+        "optimizer_lr": 0.0001,
+        "optimizer_betas": [
+            0.95,
+            0.999
+        ],
+        "optimizer_eps": 1e-08,
+        "optimizer_weight_decay": 1e-06,
+        "optimizer_vqvae_lr": 0.001,
+        "optimizer_vqvae_weight_decay": 0.0001,
+        "scheduler_warmup_steps": 500
+    },
+    "output_dir": "outputs/train/2025-01-25/20-30-40_pusht_vqbet",
+    "job_name": "pusht_vqbet",
+    "resume": false,
+    "device": "cuda",
+    "use_amp": false,
+    "seed": 100000,
+    "num_workers": 4,
+    "batch_size": 64,
+    "eval_freq": 25000,
+    "log_freq": 200,
+    "save_checkpoint": true,
+    "save_freq": 25000,
+    "offline": {
+        "steps": 250000
+    },
+    "online": {
+        "steps": 0,
+        "rollout_n_episodes": 1,
+        "rollout_batch_size": 1,
+        "steps_between_rollouts": null,
+        "sampling_ratio": 0.5,
+        "env_seed": null,
+        "buffer_capacity": null,
+        "buffer_seed_size": 0,
+        "do_rollout_async": false
+    },
+    "use_policy_training_preset": true,
+    "optimizer": {
+        "type": "adam",
+        "lr": 0.0001,
+        "betas": [
+            0.95,
+            0.999
+        ],
+        "eps": 1e-08,
+        "weight_decay": 1e-06,
+        "grad_clip_norm": 10.0
+    },
+    "scheduler": {
+        "type": "vqbet",
+        "num_warmup_steps": 500,
+        "num_vqvae_training_steps": 20000,
+        "num_cycles": 0.5
+    },
+    "eval": {
+        "n_episodes": 50,
+        "batch_size": 50,
+        "use_async_envs": false
+    },
+    "wandb": {
+        "enable": true,
+        "disable_artifact": false,
+        "project": "lerobot",
+        "entity": null,
+        "notes": null
+    }
+}