From b536510bb64c601a98bdb1cdd81fd8674113a4c6 Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 08:49:56 +0200 Subject: [PATCH 1/6] fixes mbpo --- ss2r/algorithms/mbpo/losses.py | 6 +++--- ss2r/algorithms/mbpo/model_env.py | 16 ++++++++++------ ss2r/algorithms/mbpo/on_policy_training_step.py | 10 ++++++---- 3 files changed, 19 insertions(+), 13 deletions(-) diff --git a/ss2r/algorithms/mbpo/losses.py b/ss2r/algorithms/mbpo/losses.py index de95f5ade..0c0db06df 100644 --- a/ss2r/algorithms/mbpo/losses.py +++ b/ss2r/algorithms/mbpo/losses.py @@ -83,7 +83,8 @@ def critic_loss( action = transitions.action scale = cost_scaling if safe else reward_scaling gamma = safety_discounting if safe else discounting - q_old_action = qr_network.apply( + q_net = qc_network if (safe and qc_network is not None) else qr_network + q_old_action = q_net.apply( normalizer_params, q_params, transitions.observation, action ) key, another_key = jax.random.split(key) @@ -101,7 +102,7 @@ def policy(obs: jax.Array) -> tuple[jax.Array, jax.Array]: next_action = parametric_action_distribution.postprocess(next_action) return next_action, next_log_prob - q_fn = lambda obs, action: qr_network.apply( + q_fn = lambda obs, action: q_net.apply( normalizer_params, target_q_params, obs, action ) target_q = target_q_fn( @@ -172,7 +173,6 @@ def compute_model_loss(model_params, normalizer_params, data, obs_key="state"): (next_obs_pred, reward_pred, cost_pred) = model_apply( normalizer_params, model_params, data.observation, data.action ) - next_obs_pred = normalize_fn(next_obs_pred, normalizer_params) next_obs_pred = ( next_obs_pred if isinstance(next_obs_pred, jax.Array) diff --git a/ss2r/algorithms/mbpo/model_env.py b/ss2r/algorithms/mbpo/model_env.py index 7dd689229..d2b36670a 100644 --- a/ss2r/algorithms/mbpo/model_env.py +++ b/ss2r/algorithms/mbpo/model_env.py @@ -112,9 +112,11 @@ def step(self, state: base.State, action: jax.Array) -> base.State: ).mean(axis=-1) backup_policy = self.policy_network.apply backup_policy_params = self.backup_policy_params - backup_action = backup_policy( - self.initial_normalizer_params, backup_policy_params, state.obs - )[: self.action_size] + backup_action = jnp.tanh( + backup_policy( + self.initial_normalizer_params, backup_policy_params, state.obs + )[: self.action_size] + ) qc_backup = self.qc_network.apply( self.normalizer_params, self.backup_qc_params, @@ -130,9 +132,11 @@ def step(self, state: base.State, action: jax.Array) -> base.State: pred_backup_action = self.policy_network.apply backup_policy_params = self.backup_policy_params - backup_action = pred_backup_action( - self.normalizer_params, backup_policy_params, state.obs - )[: self.action_size] + backup_action = jnp.tanh( + pred_backup_action( + self.normalizer_params, backup_policy_params, state.obs + )[: self.action_size] + ) pred_qr = self.qr_network.apply backup_qr_params = self.backup_qr_params pessimistic_qr_pred = pred_qr( diff --git a/ss2r/algorithms/mbpo/on_policy_training_step.py b/ss2r/algorithms/mbpo/on_policy_training_step.py index 5bd20ddfc..bbb83f4ce 100644 --- a/ss2r/algorithms/mbpo/on_policy_training_step.py +++ b/ss2r/algorithms/mbpo/on_policy_training_step.py @@ -179,7 +179,7 @@ def actor_sgd_step( key_alpha, optimizer_state=training_state.alpha_optimizer_state, ) - alpha = jnp.exp(training_state.alpha_params) + min_alpha + alpha = jnp.exp(alpha_params) + min_alpha (actor_loss, aux), policy_params, policy_optimizer_state = actor_update( training_state.behavior_policy_params, training_state.normalizer_params, @@ -303,9 +303,11 @@ def relabel_transitions( pred_backup_action = planning_env.policy_network.apply backup_policy_params = planning_env.backup_policy_params - backup_action = pred_backup_action( - normalizer_params, backup_policy_params, transitions.observation - )[..., : planning_env.action_size] + backup_action = jnp.tanh( + pred_backup_action( + normalizer_params, backup_policy_params, transitions.observation + )[..., : planning_env.action_size] + ) disagreement = ( next_obs_pred.std(axis=0).mean(-1) if isinstance(next_obs_pred, jax.Array) From 81288cfb18bee1e34612eb9b553c0974dda472ac Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 09:11:57 +0200 Subject: [PATCH 2/6] only fix model --- ss2r/algorithms/mbpo/losses.py | 5 ++- ss2r/algorithms/mbpo/model_env.py | 36 +++++++------------ .../mbpo/on_policy_training_step.py | 10 +++--- 3 files changed, 18 insertions(+), 33 deletions(-) diff --git a/ss2r/algorithms/mbpo/losses.py b/ss2r/algorithms/mbpo/losses.py index 0c0db06df..819bcc171 100644 --- a/ss2r/algorithms/mbpo/losses.py +++ b/ss2r/algorithms/mbpo/losses.py @@ -83,8 +83,7 @@ def critic_loss( action = transitions.action scale = cost_scaling if safe else reward_scaling gamma = safety_discounting if safe else discounting - q_net = qc_network if (safe and qc_network is not None) else qr_network - q_old_action = q_net.apply( + q_old_action = qr_network.apply( normalizer_params, q_params, transitions.observation, action ) key, another_key = jax.random.split(key) @@ -102,7 +101,7 @@ def policy(obs: jax.Array) -> tuple[jax.Array, jax.Array]: next_action = parametric_action_distribution.postprocess(next_action) return next_action, next_log_prob - q_fn = lambda obs, action: q_net.apply( + q_fn = lambda obs, action: qr_network.apply( normalizer_params, target_q_params, obs, action ) target_q = target_q_fn( diff --git a/ss2r/algorithms/mbpo/model_env.py b/ss2r/algorithms/mbpo/model_env.py index d2b36670a..6dc76a2cb 100644 --- a/ss2r/algorithms/mbpo/model_env.py +++ b/ss2r/algorithms/mbpo/model_env.py @@ -112,11 +112,9 @@ def step(self, state: base.State, action: jax.Array) -> base.State: ).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 = backup_policy( + self.initial_normalizer_params, backup_policy_params, state.obs + )[: self.action_size] qc_backup = self.qc_network.apply( self.normalizer_params, self.backup_qc_params, @@ -132,11 +130,9 @@ def step(self, state: base.State, action: jax.Array) -> base.State: 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 = pred_backup_action( + self.normalizer_params, backup_policy_params, state.obs + )[: self.action_size] pred_qr = self.qr_network.apply backup_qr_params = self.backup_qr_params pessimistic_qr_pred = pred_qr( @@ -230,13 +226,13 @@ def _propagate_ensemble( """Propagate the ensemble predictions based on the selection method.""" # Calculate the nominal predictions if ensemble_selection == "nominal": - # Get the average model parameters - avg_model_params = jax.tree_util.tree_map( - lambda p: jnp.mean(p, axis=0), model_params - ) - next_obs, reward, cost = pred_fn( - normalizer_params, avg_model_params, obs, action + vmap_pred_fn = jax.vmap(pred_fn, in_axes=(None, 0, None, None)) + next_obs_pred, reward_pred, cost_pred = vmap_pred_fn( + normalizer_params, model_params, obs, action ) + next_obs = jax.tree_map(lambda x: jnp.mean(x, axis=0), next_obs_pred) + reward = jnp.mean(reward_pred, axis=0) + cost = jnp.mean(cost_pred, axis=0) elif ensemble_selection == "random": vmap_pred_fn = jax.vmap(pred_fn, in_axes=(None, 0, None, None)) next_obs_pred, reward_pred, cost_pred = vmap_pred_fn( @@ -247,14 +243,6 @@ def _propagate_ensemble( next_obs = jax.tree_map(lambda x: x[idx], next_obs_pred) reward = reward_pred[idx] cost = cost_pred[idx] - elif ensemble_selection == "mean": - vmap_pred_fn = jax.vmap(pred_fn, in_axes=(None, 0, None, None)) - next_obs_pred, reward_pred, cost_pred = vmap_pred_fn( - normalizer_params, model_params, obs, action - ) - next_obs = jax.tree_map(lambda x: jnp.mean(x, axis=0), next_obs_pred) - reward = jnp.mean(reward_pred, axis=0) - cost = jnp.mean(cost_pred, axis=0) else: raise ValueError(f"Unknown ensemble selection: {ensemble_selection}") return next_obs, reward, cost diff --git a/ss2r/algorithms/mbpo/on_policy_training_step.py b/ss2r/algorithms/mbpo/on_policy_training_step.py index bbb83f4ce..5bd20ddfc 100644 --- a/ss2r/algorithms/mbpo/on_policy_training_step.py +++ b/ss2r/algorithms/mbpo/on_policy_training_step.py @@ -179,7 +179,7 @@ def actor_sgd_step( key_alpha, optimizer_state=training_state.alpha_optimizer_state, ) - alpha = jnp.exp(alpha_params) + min_alpha + alpha = jnp.exp(training_state.alpha_params) + min_alpha (actor_loss, aux), policy_params, policy_optimizer_state = actor_update( training_state.behavior_policy_params, training_state.normalizer_params, @@ -303,11 +303,9 @@ def relabel_transitions( 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 = pred_backup_action( + normalizer_params, backup_policy_params, transitions.observation + )[..., : planning_env.action_size] disagreement = ( next_obs_pred.std(axis=0).mean(-1) if isinstance(next_obs_pred, jax.Array) From ca2101337d55f5f875a7f969f854d1d20080e9dd Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 09:25:28 +0200 Subject: [PATCH 3/6] undo normalization bug --- ss2r/algorithms/mbpo/losses.py | 1 + ss2r/algorithms/mbpo/model_env.py | 16 ++++++++++------ ss2r/algorithms/mbpo/on_policy_training_step.py | 8 +++++--- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/ss2r/algorithms/mbpo/losses.py b/ss2r/algorithms/mbpo/losses.py index 819bcc171..de95f5ade 100644 --- a/ss2r/algorithms/mbpo/losses.py +++ b/ss2r/algorithms/mbpo/losses.py @@ -172,6 +172,7 @@ def compute_model_loss(model_params, normalizer_params, data, obs_key="state"): (next_obs_pred, reward_pred, cost_pred) = model_apply( normalizer_params, model_params, data.observation, data.action ) + next_obs_pred = normalize_fn(next_obs_pred, normalizer_params) next_obs_pred = ( next_obs_pred if isinstance(next_obs_pred, jax.Array) diff --git a/ss2r/algorithms/mbpo/model_env.py b/ss2r/algorithms/mbpo/model_env.py index 6dc76a2cb..5834dfb15 100644 --- a/ss2r/algorithms/mbpo/model_env.py +++ b/ss2r/algorithms/mbpo/model_env.py @@ -112,9 +112,11 @@ def step(self, state: base.State, action: jax.Array) -> base.State: ).mean(axis=-1) backup_policy = self.policy_network.apply backup_policy_params = self.backup_policy_params - backup_action = backup_policy( - self.initial_normalizer_params, backup_policy_params, state.obs - )[: self.action_size] + backup_action = jnp.tanh( + backup_policy( + self.initial_normalizer_params, backup_policy_params, state.obs + )[: self.action_size] + ) qc_backup = self.qc_network.apply( self.normalizer_params, self.backup_qc_params, @@ -130,9 +132,11 @@ def step(self, state: base.State, action: jax.Array) -> base.State: pred_backup_action = self.policy_network.apply backup_policy_params = self.backup_policy_params - backup_action = pred_backup_action( - self.normalizer_params, backup_policy_params, state.obs - )[: self.action_size] + backup_action = jnp.tanh( + pred_backup_action( + self.normalizer_params, backup_policy_params, state.obs + )[: self.action_size] + ) pred_qr = self.qr_network.apply backup_qr_params = self.backup_qr_params pessimistic_qr_pred = pred_qr( diff --git a/ss2r/algorithms/mbpo/on_policy_training_step.py b/ss2r/algorithms/mbpo/on_policy_training_step.py index 5bd20ddfc..c90d552a1 100644 --- a/ss2r/algorithms/mbpo/on_policy_training_step.py +++ b/ss2r/algorithms/mbpo/on_policy_training_step.py @@ -303,9 +303,11 @@ def relabel_transitions( pred_backup_action = planning_env.policy_network.apply backup_policy_params = planning_env.backup_policy_params - backup_action = pred_backup_action( - normalizer_params, backup_policy_params, transitions.observation - )[..., : planning_env.action_size] + backup_action = jnp.tanh( + pred_backup_action( + normalizer_params, backup_policy_params, transitions.observation + )[..., : planning_env.action_size] + ) disagreement = ( next_obs_pred.std(axis=0).mean(-1) if isinstance(next_obs_pred, jax.Array) From 76b0cfb5f5c8de54f76cf5de40f4d3d88f9cbb7f Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 09:39:13 +0200 Subject: [PATCH 4/6] fix cost-critic softplus bug --- ss2r/algorithms/mbpo/losses.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ss2r/algorithms/mbpo/losses.py b/ss2r/algorithms/mbpo/losses.py index de95f5ade..99b586e1c 100644 --- a/ss2r/algorithms/mbpo/losses.py +++ b/ss2r/algorithms/mbpo/losses.py @@ -83,7 +83,8 @@ def critic_loss( action = transitions.action scale = cost_scaling if safe else reward_scaling gamma = safety_discounting if safe else discounting - q_old_action = qr_network.apply( + q_net = qc_network if (safe and qc_network is not None) else qr_network + q_old_action = q_net.apply( normalizer_params, q_params, transitions.observation, action ) key, another_key = jax.random.split(key) @@ -101,7 +102,7 @@ def policy(obs: jax.Array) -> tuple[jax.Array, jax.Array]: next_action = parametric_action_distribution.postprocess(next_action) return next_action, next_log_prob - q_fn = lambda obs, action: qr_network.apply( + q_fn = lambda obs, action: q_net.apply( normalizer_params, target_q_params, obs, action ) target_q = target_q_fn( From 1f14ed155bf252800ac4d14cdc8b20f7e02621b2 Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 09:40:00 +0200 Subject: [PATCH 5/6] fix alpha params --- ss2r/algorithms/mbpo/on_policy_training_step.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ss2r/algorithms/mbpo/on_policy_training_step.py b/ss2r/algorithms/mbpo/on_policy_training_step.py index c90d552a1..bbb83f4ce 100644 --- a/ss2r/algorithms/mbpo/on_policy_training_step.py +++ b/ss2r/algorithms/mbpo/on_policy_training_step.py @@ -179,7 +179,7 @@ def actor_sgd_step( key_alpha, optimizer_state=training_state.alpha_optimizer_state, ) - alpha = jnp.exp(training_state.alpha_params) + min_alpha + alpha = jnp.exp(alpha_params) + min_alpha (actor_loss, aux), policy_params, policy_optimizer_state = actor_update( training_state.behavior_policy_params, training_state.normalizer_params, From f74ca60e0f67ccc9e6a39730a0b482d5c9d0c690 Mon Sep 17 00:00:00 2001 From: ManuelWendl Date: Tue, 19 May 2026 09:49:23 +0200 Subject: [PATCH 6/6] revert alpha change -- more stability --- ss2r/algorithms/mbpo/on_policy_training_step.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ss2r/algorithms/mbpo/on_policy_training_step.py b/ss2r/algorithms/mbpo/on_policy_training_step.py index bbb83f4ce..c90d552a1 100644 --- a/ss2r/algorithms/mbpo/on_policy_training_step.py +++ b/ss2r/algorithms/mbpo/on_policy_training_step.py @@ -179,7 +179,7 @@ def actor_sgd_step( key_alpha, optimizer_state=training_state.alpha_optimizer_state, ) - alpha = jnp.exp(alpha_params) + min_alpha + alpha = jnp.exp(training_state.alpha_params) + min_alpha (actor_loss, aux), policy_params, policy_optimizer_state = actor_update( training_state.behavior_policy_params, training_state.normalizer_params,