squash commits

cfg/robomimic/finetune/can/ft_awr_diffusion_mlp.yaml

@@ -94,13 +94,12 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_dipo_diffusion_mlp.yaml

@@ -95,13 +95,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_dql_diffusion_mlp.yaml

@@ -94,13 +94,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_idql_diffusion_mlp.yaml

@@ -97,19 +97,18 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   critic_v:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_diffusion_mlp.yaml

@@ -100,15 +100,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_diffusion_mlp_img.yaml

@@ -20,6 +20,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 ft_denoising_steps: 5
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 4
 act_steps: 4
 use_ddim: True
@@ -121,6 +122,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -133,6 +135,7 @@ model:
     time_dim: 32
     mlp_dims: [512, 512, 512]
     residual_style: True
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -143,6 +146,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -150,15 +154,14 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_diffusion_unet.yaml

@@ -103,15 +103,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_exact_diffusion_mlp.yaml

@@ -108,15 +108,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_gaussian_mlp.yaml

@@ -94,10 +94,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_gaussian_mlp_img.yaml

@@ -18,6 +18,7 @@ obs_dim: 9
 action_dim: 7
 transition_dim: ${action_dim}
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 4
 act_steps: 4
 
@@ -100,6 +101,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -115,6 +117,7 @@ model:
     learn_fixed_std: True
     std_min: 0.01
     std_max: 0.2
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -125,6 +128,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -132,10 +136,10 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_gaussian_transformer.yaml

@@ -26,7 +26,7 @@ env:
   name: ${env_name}
   best_reward_threshold_for_success: 1
   max_episode_steps: 300
-  save_video: false
+  save_video: False
   wrappers:
     robomimic_lowdim:
       normalization_path: ${normalization_path}
@@ -95,10 +95,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_gmm_mlp.yaml

@@ -27,7 +27,7 @@ env:
   name: ${env_name}
   best_reward_threshold_for_success: 1
   max_episode_steps: 300
-  save_video: false
+  save_video: False
   wrappers:
     robomimic_lowdim:
       normalization_path: ${normalization_path}
@@ -95,10 +95,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_ppo_gmm_transformer.yaml

@@ -27,7 +27,7 @@ env:
   name: ${env_name}
   best_reward_threshold_for_success: 1
   max_episode_steps: 300
-  save_video: false
+  save_video: False
   wrappers:
     robomimic_lowdim:
       normalization_path: ${normalization_path}
@@ -96,10 +96,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_qsm_diffusion_mlp.yaml

@@ -95,13 +95,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/can/ft_rwr_diffusion_mlp.yaml

@@ -86,6 +86,5 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_awr_diffusion_mlp.yaml

@@ -94,13 +94,12 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_dipo_diffusion_mlp.yaml

@@ -95,13 +95,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_dql_diffusion_mlp.yaml

@@ -94,13 +94,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_idql_diffusion_mlp.yaml

@@ -97,19 +97,18 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   critic_v:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_diffusion_mlp.yaml

@@ -100,15 +100,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_diffusion_mlp_img.yaml

@@ -20,6 +20,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 ft_denoising_steps: 5
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 4
 act_steps: 4
 use_ddim: True
@@ -121,6 +122,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -133,6 +135,7 @@ model:
     time_dim: 32
     mlp_dims: [512, 512, 512]
     residual_style: True
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -143,6 +146,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -150,15 +154,14 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_diffusion_unet.yaml

@@ -103,15 +103,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_gaussian_mlp.yaml

@@ -94,10 +94,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_gaussian_mlp_img.yaml

@@ -18,6 +18,7 @@ obs_dim: 9
 action_dim: 7
 transition_dim: ${action_dim}
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 4
 act_steps: 4
 
@@ -100,6 +101,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -115,6 +117,7 @@ model:
     learn_fixed_std: True
     std_min: 0.01
     std_max: 0.2
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -125,6 +128,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -132,10 +136,10 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_gaussian_transformer.yaml

@@ -95,10 +95,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_gmm_mlp.yaml

@@ -95,10 +95,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_ppo_gmm_transformer.yaml

@@ -96,10 +96,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_qsm_diffusion_mlp.yaml

@@ -95,13 +95,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/lift/ft_rwr_diffusion_mlp.yaml

@@ -86,6 +86,5 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_awr_diffusion_mlp.yaml

@@ -95,13 +95,12 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_dipo_diffusion_mlp.yaml

@@ -96,13 +96,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_dql_diffusion_mlp.yaml

@@ -95,13 +95,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_idql_diffusion_mlp.yaml

@@ -98,19 +98,18 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   critic_v:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_diffusion_mlp.yaml

@@ -101,15 +101,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_diffusion_mlp_img.yaml

@@ -20,6 +20,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 ft_denoising_steps: 5
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 4
 act_steps: 4
 use_ddim: True
@@ -121,6 +122,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -133,6 +135,7 @@ model:
     time_dim: 32
     mlp_dims: [768, 768, 768]
     residual_style: True
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -143,6 +146,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -150,15 +154,14 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_diffusion_unet.yaml

@@ -103,15 +103,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_gaussian_mlp.yaml

@@ -94,10 +94,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_gaussian_mlp_img.yaml

@@ -18,6 +18,7 @@ obs_dim: 9
 action_dim: 7
 transition_dim: ${action_dim}
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 4
 act_steps: 4
 
@@ -100,6 +101,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -115,6 +117,7 @@ model:
     learn_fixed_std: True
     std_min: 0.01
     std_max: 0.2
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -125,6 +128,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -132,10 +136,10 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_gaussian_transformer.yaml

@@ -95,10 +95,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_gmm_mlp.yaml

@@ -95,10 +95,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_ppo_gmm_transformer.yaml

@@ -96,10 +96,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_qsm_diffusion_mlp.yaml

@@ -96,13 +96,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/square/ft_rwr_diffusion_mlp.yaml

@@ -87,6 +87,5 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_awr_diffusion_mlp.yaml

@@ -97,13 +97,12 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_dipo_diffusion_mlp.yaml

@@ -98,13 +98,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_dql_diffusion_mlp.yaml

@@ -97,13 +97,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_idql_diffusion_mlp.yaml

@@ -100,19 +100,18 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   critic_v:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_diffusion_mlp.yaml

@@ -103,15 +103,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_diffusion_mlp_img.yaml

@@ -20,6 +20,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 ft_denoising_steps: 5
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 8
 act_steps: 8
 use_ddim: True
@@ -125,6 +126,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -138,6 +140,7 @@ model:
     time_dim: 32
     mlp_dims: [768, 768, 768]
     residual_style: True
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -149,6 +152,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -156,15 +160,14 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_diffusion_unet.yaml

@@ -106,15 +106,13 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   ft_denoising_steps: ${ft_denoising_steps}
-  transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  cond_steps: ${cond_steps}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_gaussian_mlp.yaml

@@ -97,10 +97,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_gaussian_mlp_img.yaml

@@ -18,6 +18,7 @@ obs_dim: 18
 action_dim: 14
 transition_dim: ${action_dim}
 cond_steps: 1
+img_cond_steps: 1
 horizon_steps: 8
 act_steps: 8
 
@@ -104,6 +105,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -120,6 +122,7 @@ model:
     learn_fixed_std: True
     std_min: 0.01
     std_max: 0.2
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -131,6 +134,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -138,10 +142,10 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
-    obs_dim: ${obs_dim}
+    img_cond_steps: ${img_cond_steps}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_gaussian_transformer.yaml

@@ -98,10 +98,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_gmm_mlp.yaml

@@ -98,10 +98,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_ppo_gmm_transformer.yaml

@@ -99,10 +99,9 @@ model:
     transition_dim: ${transition_dim}
   critic:
     _target_: model.common.critic.CriticObs
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
-  cond_steps: ${cond_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_qsm_diffusion_mlp.yaml

@@ -98,13 +98,12 @@ model:
     _target_: model.common.critic.CriticObsAct
     action_dim: ${action_dim}
     action_steps: ${act_steps}
-    obs_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/finetune/transport/ft_rwr_diffusion_mlp.yaml

@@ -89,6 +89,5 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
   device: ${device}

cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml

@@ -46,15 +46,13 @@ model:
     time_dim: 16
     mlp_dims: [512, 512, 512]
     residual_style: True
-    cond_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml

@@ -18,6 +18,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 horizon_steps: 4
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -27,7 +28,7 @@ wandb:
 shape_meta:
   obs:
     rgb:
-      shape: [3, 96, 96]
+      shape: [3, 96, 96]  # not counting img_cond_steps
     state:
       shape: [9]
   action: 
@@ -55,6 +56,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -62,6 +64,7 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
+    img_cond_steps: ${img_cond_steps}
     augment: True
     spatial_emb: 128
     time_dim: 32
@@ -73,9 +76,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:
@@ -88,4 +89,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml

@@ -55,9 +55,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml

@@ -17,6 +17,7 @@ action_dim: 7
 transition_dim: ${action_dim}
 horizon_steps: 4
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -52,6 +53,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -64,6 +66,7 @@ model:
     mlp_dims: [512, 512, 512]
     residual_style: True
     fixed_std: 0.1
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -80,4 +83,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml

@@ -52,9 +52,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml

@@ -18,6 +18,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 horizon_steps: 4
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -27,7 +28,7 @@ wandb:
 shape_meta:
   obs:
     rgb:
-      shape: [3, 96, 96]
+      shape: [3, 96, 96]  # not counting img_cond_steps
     state:
       shape: [9]
   action: 
@@ -55,6 +56,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -62,6 +64,7 @@ model:
         num_heads: 4
         embed_style: embed2
         embed_norm: 0
+    img_cond_steps: ${img_cond_steps}
     augment: True
     spatial_emb: 128
     time_dim: 32
@@ -73,9 +76,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:
@@ -88,4 +89,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml

@@ -51,13 +51,11 @@ model:
     smaller_encoder: False
     cond_predict_scale: True
     transition_dim: ${transition_dim}
-    cond_dim: ${obs_dim}
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml

@@ -17,6 +17,7 @@ action_dim: 7
 transition_dim: ${action_dim}
 horizon_steps: 4
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -52,6 +53,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -64,6 +66,7 @@ model:
     mlp_dims: [512, 512, 512]
     residual_style: True
     fixed_std: 0.1
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -80,4 +83,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml

@@ -53,9 +53,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml

@@ -18,6 +18,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 horizon_steps: 4
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -27,7 +28,7 @@ wandb:
 shape_meta:
   obs:
     rgb:
-      shape: [3, 96, 96]
+      shape: [3, 96, 96]  # not counting img_cond_steps
     state:
       shape: [9]
   action: 
@@ -55,6 +56,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -67,15 +69,14 @@ model:
     time_dim: 32
     mlp_dims: [768, 768, 768]
     residual_style: True
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:
@@ -88,4 +89,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml

@@ -55,9 +55,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml

@@ -17,6 +17,7 @@ action_dim: 7
 transition_dim: ${action_dim}
 horizon_steps: 4
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -52,6 +53,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -64,6 +66,7 @@ model:
     mlp_dims: [768, 768, 768]
     residual_style: True
     fixed_std: 0.1
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -80,4 +83,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml

@@ -52,9 +52,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml

@@ -18,6 +18,7 @@ transition_dim: ${action_dim}
 denoising_steps: 100
 horizon_steps: 8
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -27,7 +28,7 @@ wandb:
 shape_meta:
   obs:
     rgb:
-      shape: [3, 96, 96]
+      shape: [3, 96, 96]  # not counting img_cond_steps
     state:
       shape: [9]
   action: 
@@ -55,6 +56,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -68,15 +70,14 @@ model:
     time_dim: 32
     mlp_dims: [768, 768, 768]
     residual_style: True
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:
@@ -89,4 +90,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}

cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml

@@ -55,9 +55,7 @@ model:
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}
-  transition_dim: ${transition_dim}
   denoising_steps: ${denoising_steps}
-  cond_steps: ${cond_steps}
   device: ${device}
 
 ema:

cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml

@@ -17,6 +17,7 @@ action_dim: 14
 transition_dim: ${action_dim}
 horizon_steps: 8
 cond_steps: 1
+img_cond_steps: 1
 
 wandb:
   entity: ${oc.env:DPPO_WANDB_ENTITY}
@@ -26,7 +27,7 @@ wandb:
 shape_meta:
   obs:
     rgb:
-      shape: [3, 96, 96]
+      shape: [3, 96, 96]  # not counting img_cond_steps
     state:
       shape: [9]
   action: 
@@ -52,6 +53,7 @@ model:
     backbone:
       _target_: model.common.vit.VitEncoder
       obs_shape: ${shape_meta.obs.rgb.shape}
+      num_channel: ${eval:'${shape_meta.obs.rgb.shape[0]} * ${img_cond_steps}'} # each image patch is history concatenated
       cfg:
         patch_size: 8
         depth: 1
@@ -65,6 +67,7 @@ model:
     mlp_dims: [768, 768, 768]
     residual_style: True
     fixed_std: 0.1
+    img_cond_steps: ${img_cond_steps}
     cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     horizon_steps: ${horizon_steps}
     transition_dim: ${transition_dim}
@@ -81,4 +84,5 @@ train_dataset:
   horizon_steps: ${horizon_steps}
   max_n_episodes: 100
   cond_steps: ${cond_steps}
+  img_cond_steps: ${img_cond_steps}
   device: ${device}