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>

model/common/mlp.py

@@ -96,6 +96,7 @@ class ResidualMLP(nn.Module):
         out_activation_type="Identity",
         use_layernorm=False,
         use_layernorm_final=False,
+        dropout=0,
     ):
         super(ResidualMLP, self).__init__()
         hidden_dim = dim_list[1]
@@ -108,6 +109,7 @@ class ResidualMLP(nn.Module):
                     hidden_dim=hidden_dim,
                     activation_type=activation_type,
                     use_layernorm=use_layernorm,
+                    dropout=dropout,
                 )
                 for _ in range(1, num_hidden_layers, 2)
             ]
@@ -129,6 +131,7 @@ class TwoLayerPreActivationResNetLinear(nn.Module):
         hidden_dim,
         activation_type="Mish",
         use_layernorm=False,
+        dropout=0,
     ):
         super().__init__()
         self.l1 = nn.Linear(hidden_dim, hidden_dim)
@@ -137,6 +140,8 @@ class TwoLayerPreActivationResNetLinear(nn.Module):
         if use_layernorm:
             self.norm1 = nn.LayerNorm(hidden_dim, eps=1e-06)
             self.norm2 = nn.LayerNorm(hidden_dim, eps=1e-06)
+        if dropout > 0:
+            raise NotImplementedError("Dropout not implemented for residual MLP!")
 
     def forward(self, x):
         x_input = x

model/common/mlp_gaussian.py

@@ -212,6 +212,7 @@ class Gaussian_MLP(nn.Module):
                 out_activation_type=activation_type,
                 use_layernorm=use_layernorm,
                 use_layernorm_final=use_layernorm,
+                dropout=dropout,
             )
             self.mlp_mean = MLP(
                 mlp_dims[-1:] + [output_dim],
@@ -233,9 +234,7 @@ class Gaussian_MLP(nn.Module):
             if learn_fixed_std:
                 # initialize to fixed_std
                 self.logvar = torch.nn.Parameter(
-                    torch.log(
-                        torch.tensor([fixed_std**2 for _ in range(action_dim)])
-                    ),
+                    torch.log(torch.tensor([fixed_std**2 for _ in range(action_dim)])),
                     requires_grad=True,
                 )
         self.logvar_min = torch.nn.Parameter(