From 78ce9ac1fa830bf340e37acd45c5bc9313908ce2 Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 10:39:14 +0200 Subject: [PATCH 1/3] fix tanh directly --- ss2r/algorithms/mbpo/model_env.py | 26 ++++++++++--------- .../mbpo/on_policy_training_step.py | 8 ++---- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/ss2r/algorithms/mbpo/model_env.py b/ss2r/algorithms/mbpo/model_env.py index 5834dfb1..c7b07aae 100644 --- a/ss2r/algorithms/mbpo/model_env.py +++ b/ss2r/algorithms/mbpo/model_env.py @@ -3,6 +3,7 @@ from brax import envs from brax.envs import base +from ss2r.algorithms.mbpo.networks import make_inference_fn from ss2r.algorithms.sac.types import float32 @@ -46,6 +47,15 @@ def __init__( self.initial_normalizer_params = ( initial_normalizer_params if initial_normalizer_params is not None else {} ) + _make_policy = make_inference_fn(mbpo_network) + self.backup_policy_fn = _make_policy( + (training_state.normalizer_params, training_state.backup_policy_params), + deterministic=True, + ) + self.backup_policy_fn_initial = _make_policy( + (self.initial_normalizer_params, training_state.backup_policy_params), + deterministic=True, + ) def reset(self, rng: jax.Array) -> base.State: sample_key, model_key = jax.random.split(rng) @@ -110,12 +120,8 @@ def step(self, state: base.State, action: jax.Array) -> base.State: state.obs, action, ).mean(axis=-1) - backup_policy = self.policy_network.apply - backup_policy_params = self.backup_policy_params - backup_action = jnp.tanh( - backup_policy( - self.initial_normalizer_params, backup_policy_params, state.obs - )[: self.action_size] + backup_action, _ = self.backup_policy_fn_initial( + state.obs, jax.random.PRNGKey(0) ) qc_backup = self.qc_network.apply( self.normalizer_params, @@ -130,12 +136,8 @@ def step(self, state: base.State, action: jax.Array) -> base.State: done, ) - pred_backup_action = self.policy_network.apply - backup_policy_params = self.backup_policy_params - backup_action = jnp.tanh( - pred_backup_action( - self.normalizer_params, backup_policy_params, state.obs - )[: self.action_size] + backup_action, _ = self.backup_policy_fn( + state.obs, jax.random.PRNGKey(0) ) pred_qr = self.qr_network.apply backup_qr_params = self.backup_qr_params diff --git a/ss2r/algorithms/mbpo/on_policy_training_step.py b/ss2r/algorithms/mbpo/on_policy_training_step.py index c90d552a..ce7ffdaa 100644 --- a/ss2r/algorithms/mbpo/on_policy_training_step.py +++ b/ss2r/algorithms/mbpo/on_policy_training_step.py @@ -301,12 +301,8 @@ def relabel_transitions( normalizer_params, model_params, transitions.observation, transitions.action ) - pred_backup_action = planning_env.policy_network.apply - backup_policy_params = planning_env.backup_policy_params - backup_action = jnp.tanh( - pred_backup_action( - normalizer_params, backup_policy_params, transitions.observation - )[..., : planning_env.action_size] + backup_action, _ = planning_env.backup_policy_fn( + transitions.observation, jax.random.PRNGKey(0) ) disagreement = ( next_obs_pred.std(axis=0).mean(-1) From 07bc13474b3f58f192ab93aba07d08a4166e28f3 Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 11:02:29 +0200 Subject: [PATCH 2/3] keep normalizer for backup policy --- ss2r/algorithms/mbpo/model_env.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/ss2r/algorithms/mbpo/model_env.py b/ss2r/algorithms/mbpo/model_env.py index c7b07aae..ddc56a47 100644 --- a/ss2r/algorithms/mbpo/model_env.py +++ b/ss2r/algorithms/mbpo/model_env.py @@ -49,10 +49,6 @@ def __init__( ) _make_policy = make_inference_fn(mbpo_network) self.backup_policy_fn = _make_policy( - (training_state.normalizer_params, training_state.backup_policy_params), - deterministic=True, - ) - self.backup_policy_fn_initial = _make_policy( (self.initial_normalizer_params, training_state.backup_policy_params), deterministic=True, ) @@ -120,7 +116,7 @@ def step(self, state: base.State, action: jax.Array) -> base.State: state.obs, action, ).mean(axis=-1) - backup_action, _ = self.backup_policy_fn_initial( + backup_action, _ = self.backup_policy_fn( state.obs, jax.random.PRNGKey(0) ) qc_backup = self.qc_network.apply( From 76f6e2e057ccd5ab254d38de24da498de56aec9a Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 11:42:40 +0200 Subject: [PATCH 3/3] fix ruff --- ss2r/algorithms/mbpo/model_env.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ss2r/algorithms/mbpo/model_env.py b/ss2r/algorithms/mbpo/model_env.py index ddc56a47..a85c0bb6 100644 --- a/ss2r/algorithms/mbpo/model_env.py +++ b/ss2r/algorithms/mbpo/model_env.py @@ -132,9 +132,7 @@ def step(self, state: base.State, action: jax.Array) -> base.State: done, ) - backup_action, _ = self.backup_policy_fn( - state.obs, jax.random.PRNGKey(0) - ) + backup_action, _ = self.backup_policy_fn(state.obs, jax.random.PRNGKey(0)) pred_qr = self.qr_network.apply backup_qr_params = self.backup_qr_params pessimistic_qr_pred = pred_qr(