Updates from the sample efficiency branch

ss2r/algorithms/sac/train.py

@@ -152,6 +152,7 @@ def train(
     reset_on_eval: bool = True,
     store_buffer: bool = False,
     use_rae: bool = False,
+    critic_entropy: bool = True,
 ):
     if min_replay_size >= num_timesteps:
         raise ValueError(
@@ -335,6 +336,7 @@ def train(
         safety_budget,
         tau,
         num_critic_updates_per_actor_update,
+        critic_entropy=critic_entropy,
     )
 
     def prefill_replay_buffer(

ss2r/algorithms/sac/training_step.py

@@ -34,6 +34,7 @@ def make_training_step(
     safety_budget,
     tau,
     num_critic_updates_per_actor_update,
+    critic_entropy=True,
 ):
     def critic_sgd_step(
         carry: Tuple[TrainingState, PRNGKey], transitions: Transition
@@ -42,7 +43,9 @@ def critic_sgd_step(
 
         key, key_critic, key_cost_critic = jax.random.split(key, 3)
         transitions = float32(transitions)
-        alpha = jnp.exp(training_state.alpha_params) + min_alpha
+        alpha = (
+            jnp.exp(training_state.alpha_params) + min_alpha if critic_entropy else 0.0
+        )
         critic_loss, qr_params, qr_optimizer_state = critic_update(
             training_state.qr_params,
             training_state.policy_params,

ss2r/benchmark_suites/wrappers.py

@@ -448,5 +448,5 @@ def where_done(x, y):
         else:
             raise NotImplementedError
         new_data = jax.tree.map(where_done, maybe_reset_data, state_data)
-        obs = jax.tree.map(maybe_reset.obs, state.obs)
+        obs = jax.tree.map(where_done, maybe_reset.obs, state.obs)
         return state.replace(**{data_name: new_data, "obs": obs})

ss2r/configs/agent/sac.yaml

@@ -34,3 +34,4 @@ store_buffer: false
 use_rae: false
 n_critics: 2
 n_heads: 1
+critic_entropy: true

ss2r/configs/experiment/go1_sim_to_real_rae.yaml

@@ -9,9 +9,11 @@ training:
   train_domain_randomization: true
   eval_domain_randomization: true
   safe: false
-  wandb_id: 9a3kb4n0
-  num_envs: 1
-  num_timesteps: 5000000
+  wandb_id: 8wuunrrv
+  num_envs: 2
+  num_timesteps: 1000000
+  num_evals: 20
+  hard_resets: true
 
 environment:
   train_params:
@@ -44,3 +46,4 @@ agent:
   num_critic_updates_per_actor_update: 20
   n_heads: 10
   n_critics: 1
+  critic_entropy: false