v0.7 (#26)

* update from scratch configs

* update gym pretraining configs - use fewer epochs

* update robomimic pretraining configs - use fewer epochs

* allow trajectory plotting in eval agent

* add simple vit unet

* update avoid pretraining configs - use fewer epochs

* update furniture pretraining configs - use same amount of epochs as before

* add robomimic diffusion unet pretraining configs

* update robomimic finetuning configs - higher lr

* add vit unet checkpoint urls

* update pretraining and finetuning instructions as configs are updated

cfg/gym/scratch/walker2d-v2/ppo_diffusion_mlp.yaml

@@ -42,7 +42,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 501
   n_critic_warmup_itr: 0
   n_steps: 1000
   gamma: 0.99
@@ -55,7 +55,7 @@ train:
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
-    first_cycle_steps: 10000
+    first_cycle_steps: 1000
     warmup_steps: 10
     min_lr: 1e-3
   save_model_freq: 100
@@ -67,7 +67,7 @@ train:
   reward_scale_running: True
   reward_scale_const: 1.0
   gae_lambda: 0.95
-  batch_size: 10000
+  batch_size: 5000
   update_epochs: 10
   vf_coef: 0.5
   target_kl: 1
@@ -75,7 +75,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_ppo.PPODiffusion
   # HP to tune
-  gamma_denoising: 0.99
+  gamma_denoising: 1
   clip_ploss_coef: 0.1
   clip_ploss_coef_base: 0.1
   clip_ploss_coef_rate: 3
@@ -94,10 +94,10 @@ model:
     residual_style: True
   critic:
     _target_: model.common.critic.CriticObs
-    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
     mlp_dims: [256, 256, 256]
     activation_type: Mish
     residual_style: True
+    cond_dim: ${eval:'${obs_dim} * ${cond_steps}'}
   ft_denoising_steps: ${ft_denoising_steps}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}

cfg/gym/scratch/walker2d-v2/ppo_gaussian_mlp.yaml

@@ -40,7 +40,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 301
   n_critic_warmup_itr: 0
   n_steps: 1000
   gamma: 0.99
@@ -65,7 +65,7 @@ train:
   reward_scale_running: True
   reward_scale_const: 1.0
   gae_lambda: 0.95
-  batch_size: 1000
+  batch_size: 500
   update_epochs: 10
   vf_coef: 0.5
   target_kl: 1