v0.6 (#18)

* Sampling over both env and denoising steps in DPPO updates (#13)

* sample one from each chain

* full random sampling

* Add Proficient Human (PH) Configs and Pipeline (#16)

* fix missing cfg

* add ph config

* fix how terminated flags are added to buffer in ibrl

* add ph config

* offline calql for 1M gradient updates

* bug fix: number of calql online gradient steps is the number of new transitions collected

* add sample config for DPPO with ta=1

* Sampling over both env and denoising steps in DPPO updates (#13)

* sample one from each chain

* full random sampling

* fix diffusion loss when predicting initial noise

* fix dppo inds

* fix typo

* remove print statement

---------

Co-authored-by: Justin M. Lidard <jlidard@neuronic.cs.princeton.edu>
Co-authored-by: allenzren <allen.ren@princeton.edu>

* update robomimic configs

* better calql formulation

* optimize calql and ibrl training

* optimize data transfer in ppo agents

* add kitchen configs

* re-organize config folders, rerun calql and rlpd

* add scratch gym locomotion configs

* add kitchen installation dependencies

* use truncated for termination in furniture env

* update furniture and gym configs

* update README and dependencies with kitchen

* add url for new data and checkpoints

* update demo RL configs

* update batch sizes for furniture unet configs

* raise error about dropout in residual mlp

* fix observation bug in bc loss

---------

Co-authored-by: Justin Lidard <60638575+jlidard@users.noreply.github.com>
Co-authored-by: Justin M. Lidard <jlidard@neuronic.cs.princeton.edu>

cfg/robomimic/finetune/transport/ft_awr_diffusion_mlp.yaml

@@ -26,7 +26,7 @@ env:
   name: ${env_name}
   best_reward_threshold_for_success: 1
   max_episode_steps: 800
-  save_video: false
+  save_video: False
   wrappers:
     robomimic_lowdim:
       normalization_path: ${normalization_path}
@@ -49,7 +49,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 201
   n_critic_warmup_itr: 2
   n_steps: 400
   gamma: 0.999
@@ -58,7 +58,7 @@ train:
   actor_lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
@@ -82,7 +82,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_awr.AWRDiffusion
   # Sampling HPs
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   randn_clip_value: 3
   #
   network_path: ${base_policy_path}

cfg/robomimic/finetune/transport/ft_dipo_diffusion_mlp.yaml

@@ -49,7 +49,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 201
   n_critic_warmup_itr: 2
   n_steps: 400
   gamma: 0.999
@@ -58,7 +58,7 @@ train:
   actor_lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
@@ -82,7 +82,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_dipo.DIPODiffusion
   # HP to tune
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   randn_clip_value: 3
   #
   network_path: ${base_policy_path}
@@ -96,12 +96,12 @@ model:
     action_dim: ${action_dim}
   critic:
     _target_: model.common.critic.CriticObsAct
-    action_dim: ${action_dim}
-    action_steps: ${act_steps}
-    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}'}
+    action_dim: ${action_dim}
+    action_steps: ${act_steps}
   horizon_steps: ${horizon_steps}
   obs_dim: ${obs_dim}
   action_dim: ${action_dim}

cfg/robomimic/finetune/transport/ft_dql_diffusion_mlp.yaml

@@ -26,7 +26,7 @@ env:
   name: ${env_name}
   best_reward_threshold_for_success: 1
   max_episode_steps: 800
-  save_video: false
+  save_video: False
   wrappers:
     robomimic_lowdim:
       normalization_path: ${normalization_path}
@@ -49,8 +49,8 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
-  n_critic_warmup_itr: 2
+  n_train_itr: 201
+  n_critic_warmup_itr: 5
   n_steps: 400
   gamma: 0.999
   actor_lr: 1e-5
@@ -58,7 +58,7 @@ train:
   actor_lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
@@ -81,7 +81,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_dql.DQLDiffusion
   # Sampling HPs
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   randn_clip_value: 3
   #
   network_path: ${base_policy_path}

cfg/robomimic/finetune/transport/ft_idql_diffusion_mlp.yaml

@@ -49,7 +49,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 201
   n_critic_warmup_itr: 5
   n_steps: 400
   gamma: 0.999
@@ -58,7 +58,7 @@ train:
   actor_lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
@@ -83,7 +83,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_idql.IDQLDiffusion
   # Sampling HPs
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   randn_clip_value: 3
   #
   network_path: ${base_policy_path}

cfg/robomimic/finetune/transport/ft_ppo_diffusion_mlp.yaml

@@ -50,16 +50,16 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 201
   n_critic_warmup_itr: 2
   n_steps: 400
   gamma: 0.999
-  actor_lr: 1e-5
+  actor_lr: 1e-4
   actor_weight_decay: 0
   actor_lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-4
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
@@ -76,7 +76,7 @@ train:
   reward_scale_const: 1.0
   gae_lambda: 0.95
   batch_size: 10000
-  update_epochs: 8
+  update_epochs: 5
   vf_coef: 0.5
   target_kl: 1
 
@@ -88,7 +88,7 @@ model:
   clip_ploss_coef_base: 0.001
   clip_ploss_coef_rate: 3
   randn_clip_value: 3
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   min_logprob_denoising_std: 0.1
   #
   network_path: ${base_policy_path}
@@ -102,10 +102,10 @@ model:
     action_dim: ${action_dim}
   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/robomimic/finetune/transport/ft_ppo_diffusion_mlp_img.yaml

@@ -1,7 +1,7 @@
 defaults:
   - _self_
 hydra:
-  run:  
+  run:
     dir: ${logdir}
 _target_: agent.finetune.train_ppo_diffusion_img_agent.TrainPPOImgDiffusionAgent
 
@@ -64,7 +64,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 500
+  n_train_itr: 201
   n_critic_warmup_itr: 2
   n_steps: 400
   gamma: 0.999
@@ -73,13 +73,13 @@ train:
   actor_lr: 1e-5
   actor_weight_decay: 0
   actor_lr_scheduler:
-    first_cycle_steps: 500
+    first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
-    first_cycle_steps: 500
+    first_cycle_steps: 1000
     warmup_steps: 10
     min_lr: 1e-3
   save_model_freq: 100
@@ -93,19 +93,19 @@ train:
   gae_lambda: 0.95
   batch_size: 500
   logprob_batch_size: 1000
-  update_epochs: 8
+  update_epochs: 10
   vf_coef: 0.5
   target_kl: 1
 
 model:
   _target_: model.diffusion.diffusion_ppo.PPODiffusion
   # HP to tune
-  gamma_denoising: 0.9
+  gamma_denoising: 0.99
   clip_ploss_coef: 0.01
   clip_ploss_coef_base: 0.001
   clip_ploss_coef_rate: 3
   randn_clip_value: 3
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   min_logprob_denoising_std: 0.1
   #
   use_ddim: ${use_ddim}
@@ -164,10 +164,10 @@ model:
         embed_style: embed2
         embed_norm: 0
     img_cond_steps: ${img_cond_steps}
-    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/robomimic/finetune/transport/ft_qsm_diffusion_mlp.yaml

@@ -49,7 +49,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 201
   n_critic_warmup_itr: 5
   n_steps: 400
   gamma: 0.999
@@ -58,7 +58,7 @@ train:
   actor_lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   critic_lr: 1e-3
   critic_weight_decay: 0
   critic_lr_scheduler:
@@ -81,7 +81,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_qsm.QSMDiffusion
   # Sampling HPs
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   randn_clip_value: 3
   #
   network_path: ${base_policy_path}

cfg/robomimic/finetune/transport/ft_rwr_diffusion_mlp.yaml

@@ -49,7 +49,7 @@ wandb:
   run: ${now:%H-%M-%S}_${name}
 
 train:
-  n_train_itr: 1000
+  n_train_itr: 201
   n_critic_warmup_itr: 2
   n_steps: 400
   gamma: 0.999
@@ -58,7 +58,7 @@ train:
   lr_scheduler:
     first_cycle_steps: 1000
     warmup_steps: 10
-    min_lr: 1e-6
+    min_lr: 1e-5
   save_model_freq: 100
   val_freq: 10
   render:
@@ -73,7 +73,7 @@ train:
 model:
   _target_: model.diffusion.diffusion_rwr.RWRDiffusion
   # Sampling HPs
-  min_sampling_denoising_std: 0.08
+  min_sampling_denoising_std: 0.1
   randn_clip_value: 3
   #
   network_path: ${base_policy_path}