Upload 3 files

config.json

@@ -1,7 +1,7 @@
 {
   "n_obs_steps": 5,
   "n_action_pred_token": 7,
-  "n_action_pred_chunk": 5,
+  "action_chunk_size": 5,
   "input_shapes": {
     "observation.image": [
       3,
@@ -33,16 +33,22 @@
   "pretrained_backbone_weights": null,
   "use_group_norm": true,
   "spatial_softmax_num_keypoints": 32,
-  "discretize_step": 20000,
+  "n_vqvae_training_steps": 20000,
   "vqvae_groups": 2,
   "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
+  "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
 }

config.yaml

@@ -1,3 +1,4 @@
+resume: false
 device: cuda
 use_amp: false
 seed: 100000
@@ -11,7 +12,8 @@ training:
   eval_freq: 20000
   save_freq: 20000
   log_freq: 250
-  save_model: true
+  save_checkpoint: true
+  num_workers: 4
   batch_size: 64
   grad_clip_norm: 10
   lr: 0.0001
@@ -23,7 +25,7 @@ training:
   adam_eps: 1.0e-08
   adam_weight_decay: 1.0e-06
   vqvae_lr: 0.001
-  discretize_step: 20000
+  n_vqvae_training_steps: 20000
   bet_weight_decay: 0.0002
   bet_learning_rate: 5.5e-05
   bet_betas:
@@ -71,13 +73,16 @@ fps: 10
 env:
   name: pusht
   task: PushT-v0
-  from_pixels: true
-  pixels_only: false
   image_size: 96
-  episode_length: 300
-  fps: ${fps}
   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:
@@ -106,7 +111,7 @@ policy:
   name: vqbet
   n_obs_steps: 5
   n_action_pred_token: 7
-  n_action_pred_chunk: 5
+  action_chunk_size: 5
   input_shapes:
     observation.image:
     - 3
@@ -130,7 +135,7 @@ policy:
   pretrained_backbone_weights: null
   use_group_norm: true
   spatial_softmax_num_keypoints: 32
-  discretize_step: ${training.discretize_step}
+  discretize_step: ${training.n_vqvae_training_steps}
   vqvae_groups: 2
   vqvae_n_embed: 16
   vqvae_embedding_dim: 256
@@ -141,10 +146,10 @@ policy:
   gpt_n_layer: 8
   gpt_n_head: 8
   gpt_hidden_dim: 512
-  gpt_num_obs_mode: 2
   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:8805c1adaf7a5680e00fa00cc286423d68cd2ec3d1fd7ffa6a8e514d5ada1beb
-size 158153964
+oid sha256:df4f83724946d68ed33205379f7d82a9fb152577c330d5de86e8b270761f37f3
+size 158154222