From 85a727c89528db37599de63c7cdee6d6d0612aa4 Mon Sep 17 00:00:00 2001
From: Yarden <yarden.as@inf.ethz.ch>
Date: Tue, 10 Jun 2025 07:44:50 +0200
Subject: [PATCH] Updates from the sample efficiency branch

---
 ss2r/algorithms/sac/train.py                     | 2 ++
 ss2r/algorithms/sac/training_step.py             | 5 ++++-
 ss2r/benchmark_suites/wrappers.py                | 2 +-
 ss2r/configs/agent/sac.yaml                      | 1 +
 ss2r/configs/experiment/go1_sim_to_real_rae.yaml | 9 ++++++---
 5 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/ss2r/algorithms/sac/train.py b/ss2r/algorithms/sac/train.py
index 8776cfa1b..39dc916de 100644
--- a/ss2r/algorithms/sac/train.py
+++ b/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(
diff --git a/ss2r/algorithms/sac/training_step.py b/ss2r/algorithms/sac/training_step.py
index a63aa2f43..89cebaec6 100644
--- a/ss2r/algorithms/sac/training_step.py
+++ b/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,
diff --git a/ss2r/benchmark_suites/wrappers.py b/ss2r/benchmark_suites/wrappers.py
index 0698f3ca6..abe019d3d 100644
--- a/ss2r/benchmark_suites/wrappers.py
+++ b/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})
diff --git a/ss2r/configs/agent/sac.yaml b/ss2r/configs/agent/sac.yaml
index 3f64640df..7f92ba2a6 100644
--- a/ss2r/configs/agent/sac.yaml
+++ b/ss2r/configs/agent/sac.yaml
@@ -34,3 +34,4 @@ store_buffer: false
 use_rae: false
 n_critics: 2
 n_heads: 1
+critic_entropy: true
\ No newline at end of file
diff --git a/ss2r/configs/experiment/go1_sim_to_real_rae.yaml b/ss2r/configs/experiment/go1_sim_to_real_rae.yaml
index faa59f6ba..58aac42ff 100644
--- a/ss2r/configs/experiment/go1_sim_to_real_rae.yaml
+++ b/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
\ No newline at end of file