Add penalizer to mbpo (#131)

* Training of behavior qc

* Rename to backup and behavior

* Don't terminate or stop rewards if penalizer is on

Author: yardenas

pyproject.toml

@@ -100,7 +100,6 @@ skip-magic-trailing-comma = false
 line-ending = "auto"
 
 [tool.mypy]
-plugins = ["numpy.typing.mypy_plugin"]
 ignore_missing_imports = true
 show_column_numbers = true
 disallow_untyped_defs = false

ss2r/algorithms/mbpo/__init__.py

@@ -1,6 +1,7 @@
 import functools
 
 import ss2r.algorithms.mbpo.networks as mbpo_networks
+from ss2r.algorithms.penalizers import get_penalizer
 from ss2r.algorithms.sac.data import get_collection_fn
 from ss2r.algorithms.sac.q_transforms import (
     get_cost_q_transform,
@@ -56,6 +57,7 @@ def get_train_fn(cfg, checkpoint_path, restore_checkpoint_path):
         value_obs_key=value_obs_key,
         policy_obs_key=policy_obs_key,
     )
+    penalizer, penalizer_params = get_penalizer(cfg)
     reward_q_transform = get_reward_q_transform(cfg)
     cost_q_transform = get_cost_q_transform(cfg)
     data_collection = get_collection_fn(cfg)
@@ -67,6 +69,8 @@ def get_train_fn(cfg, checkpoint_path, restore_checkpoint_path):
         checkpoint_logdir=checkpoint_path,
         reward_q_transform=reward_q_transform,
         cost_q_transform=cost_q_transform,
+        penalizer=penalizer,
+        penalizer_params=penalizer_params,
         get_experience_fn=data_collection,
         restore_checkpoint_path=restore_checkpoint_path,
     )

ss2r/algorithms/mbpo/losses.py

@@ -26,6 +26,7 @@
 from brax.training.types import Params, PRNGKey
 
 from ss2r.algorithms.mbpo.networks import MBPONetworks
+from ss2r.algorithms.penalizers import Penalizer
 from ss2r.algorithms.sac.q_transforms import QTransformation
 
 Transition: TypeAlias = types.Transition
@@ -45,6 +46,7 @@ def make_losses(
     target_entropy = -0.5 * action_size if init_alpha is None else init_alpha
     policy_network = mbpo_network.policy_network
     qr_network = mbpo_network.qr_network
+    qc_network = mbpo_network.qc_network
     parametric_action_distribution = mbpo_network.parametric_action_distribution
 
     def alpha_loss(
@@ -122,9 +124,13 @@ def actor_loss(
         policy_params: Params,
         normalizer_params: Any,
         qr_params: Params,
+        qc_params: Params | None,
         alpha: jnp.ndarray,
         transitions: Transition,
         key: PRNGKey,
+        safety_budget: float,
+        penalizer: Penalizer | None,
+        penalizer_params: Any,
     ) -> jnp.ndarray:
         dist_params = policy_network.apply(
             normalizer_params, policy_params, transitions.observation
@@ -143,8 +149,23 @@ def actor_loss(
             qr = jnp.min(qr_action, axis=-1)
         actor_loss = -qr.mean()
         exploration_loss = (alpha * log_prob).mean()
+        aux = {}
+        if penalizer is not None:
+            assert qc_network is not None
+            qc_action = qc_network.apply(
+                normalizer_params, qc_params, transitions.observation, action
+            )
+            mean_qc = jnp.mean(qc_action, axis=-1)
+            constraint = safety_budget - mean_qc.mean() / cost_scaling
+            actor_loss, penalizer_aux, penalizer_params = penalizer(
+                actor_loss, constraint, jax.lax.stop_gradient(penalizer_params)
+            )
+            aux["constraint_estimate"] = constraint
+            aux["cost"] = mean_qc.mean() / cost_scaling
+            aux["penalizer_params"] = penalizer_params
+            aux |= penalizer_aux
         actor_loss += exploration_loss
-        return actor_loss
+        return actor_loss, aux
 
     def compute_model_loss(model_params, normalizer_params, data, obs_key="state"):
         model_apply = jax.vmap(mbpo_network.model_network.apply, (None, 0, None, None))

ss2r/algorithms/mbpo/model_env.py

@@ -26,7 +26,7 @@ def __init__(
         self.model_network = mbpo_network.model_network
         self.model_params = training_state.model_params
         self.qc_network = mbpo_network.qc_network
-        self.qc_params = training_state.qc_params
+        self.backup_qc_params = training_state.backup_qc_params
         self.qr_network = mbpo_network.qr_network
         self.backup_qr_params = training_state.backup_qr_params
         self.policy_network = mbpo_network.policy_network
@@ -86,7 +86,7 @@ def step(self, state: base.State, action: jax.Array) -> base.State:
             expected_cost_for_traj = prev_cumulative_cost + self.scaling_fn(
                 self.qc_network.apply(
                     self.normalizer_params,
-                    self.qc_params,
+                    self.backup_qc_params,
                     state.obs,
                     action,
                 ).mean(axis=-1)

ss2r/algorithms/mbpo/safe_rollout.py

@@ -13,12 +13,12 @@ def get_inference_policy_params(safe: bool, safety_budget=float("inf")) -> Any:
     def get_params(training_state: TrainingState) -> Any:
         if safe:
             return (
-                training_state.policy_params,
-                training_state.qc_params,
+                training_state.behavior_policy_params,
+                training_state.backup_qc_params,
                 safety_budget,
             )
         else:
-            return training_state.policy_params
+            return training_state.behavior_policy_params
 
     return get_params
 

ss2r/algorithms/mbpo/train.py

@@ -41,6 +41,7 @@
 )
 from ss2r.algorithms.mbpo.training_step import make_training_step
 from ss2r.algorithms.mbpo.types import TrainingState
+from ss2r.algorithms.penalizers import Params, Penalizer
 from ss2r.algorithms.sac import gradients
 from ss2r.algorithms.sac.data import collect_single_step
 from ss2r.algorithms.sac.q_transforms import QTransformation, SACBase, SACCost
@@ -87,6 +88,7 @@ def _init_training_state(
     qc_optimizer: optax.GradientTransformation,
     model_optimizer: optax.GradientTransformation,
     model_ensemble_size: int,
+    penalizer_params: Params | None,
 ) -> TrainingState:
     """Inits the training state and replicates it over devices."""
     key_policy, key_qr, key_model = jax.random.split(key, 3)
@@ -101,16 +103,14 @@ def _init_training_state(
     model_params = init_model_ensemble(model_keys)
     model_optimizer_state = model_optimizer.init(model_params)
     if mbpo_network.qc_network is not None:
-        qc_params = mbpo_network.qc_network.init(key_qr)
+        backup_qc_params = mbpo_network.qc_network.init(key_qr)
         assert qc_optimizer is not None
-        qc_optimizer_state = qc_optimizer.init(qc_params)
+        backup_qc_optimizer_state = qc_optimizer.init(backup_qc_params)
         backup_qr_params = qr_params
-        backup_qr_optimizer_state = qr_optimizer_state
     else:
-        qc_params = None
-        qc_optimizer_state = None
+        backup_qc_params = None
+        backup_qc_optimizer_state = None
         backup_qr_params = None
-        backup_qr_optimizer_state = None
     if isinstance(obs_size, Mapping):
         obs_shape = {
             k: specs.Array(v, jnp.dtype("float32")) for k, v in obs_size.items()
@@ -119,24 +119,27 @@ def _init_training_state(
         obs_shape = specs.Array((obs_size,), jnp.dtype("float32"))
     normalizer_params = running_statistics.init_state(obs_shape)
     training_state = TrainingState(
-        policy_optimizer_state=policy_optimizer_state,
-        policy_params=policy_params,
+        behavior_policy_optimizer_state=policy_optimizer_state,
+        behavior_policy_params=policy_params,
         backup_policy_params=policy_params,
-        qr_optimizer_state=qr_optimizer_state,
-        qr_params=qr_params,
-        backup_qr_optimizer_state=backup_qr_optimizer_state,
+        behavior_qr_optimizer_state=qr_optimizer_state,
+        behavior_qr_params=qr_params,
         backup_qr_params=backup_qr_params,
-        qc_optimizer_state=qc_optimizer_state,
-        qc_params=qc_params,
-        target_qr_params=qr_params,
-        target_qc_params=qc_params,
+        behavior_qc_optimizer_state=backup_qc_optimizer_state,
+        behavior_qc_params=backup_qc_params,
+        behavior_target_qr_params=qr_params,
+        behavior_target_qc_params=backup_qc_params,
+        backup_qc_params=backup_qc_params,
+        backup_qc_optimizer_state=backup_qc_optimizer_state,
+        backup_target_qc_params=backup_qc_params,
         model_params=model_params,
         model_optimizer_state=model_optimizer_state,
         gradient_steps=jnp.zeros(()),
         env_steps=jnp.zeros(()),
         alpha_optimizer_state=alpha_optimizer_state,
         alpha_params=log_alpha,
         normalizer_params=normalizer_params,
+        penalizer_params=penalizer_params,
     )  #  type: ignore
     return training_state
 
@@ -188,6 +191,8 @@ def train(
     eval_env: Optional[envs.Env] = None,
     safe: bool = False,
     safety_budget: float = float("inf"),
+    penalizer: Penalizer | None = None,
+    penalizer_params: Params | None = None,
     reward_q_transform: QTransformation = SACBase(),
     cost_q_transform: QTransformation = SACCost(),
     use_bro: bool = True,
@@ -302,6 +307,7 @@ def train(
         qc_optimizer=qc_optimizer,
         model_optimizer=model_optimizer,
         model_ensemble_size=model_ensemble_size,
+        penalizer_params=penalizer_params,
     )
     del global_key
     local_key, model_rb_key, actor_critic_rb_key, env_key, eval_key = jax.random.split(
@@ -318,13 +324,13 @@ def train(
             ts_normalizer_params = params[0]
         training_state = training_state.replace(  # type: ignore
             normalizer_params=ts_normalizer_params,
-            policy_params=params[1],
+            behavior_policy_params=params[1],
             backup_policy_params=params[1],
-            qr_params=params[3],
+            behavior_qr_params=params[3],
             backup_qr_params=params[3],
-            qc_params=params[4] if safe else None,
+            behavior_qc_params=params[4] if safe else None,
+            backup_qc_params=params[4] if safe else None,
         )
-
     make_planning_policy = mbpo_networks.make_inference_fn(mbpo_network)
     if safe:
         make_rollout_policy = make_safe_inference_fn(
@@ -386,7 +392,7 @@ def train(
     )
     actor_update = (
         gradients.gradient_update_fn(  # pytype: disable=wrong-arg-types  # jax-ndarray
-            actor_loss, policy_optimizer, pmap_axis_name=None
+            actor_loss, policy_optimizer, pmap_axis_name=None, has_aux=True
         )
     )
     extra_fields = ("truncation",)
@@ -402,7 +408,7 @@ def train(
         safety_budget=safety_budget,
         cost_discount=safety_discounting,
         scaling_fn=budget_scaling_fun,
-        use_termination=use_termination,
+        use_termination=penalizer is not None and use_termination,
     )
     training_step = make_training_step(
         env,
@@ -434,7 +440,9 @@ def train(
         pessimism,
         model_to_real_data_ratio,
         budget_scaling_fun,
-        use_termination=use_termination,
+        use_termination,
+        penalizer,
+        safety_budget,
     )
 
     def prefill_replay_buffer(
@@ -635,9 +643,9 @@ def training_epoch_with_timing(
             # Save current policy.
             params = (
                 training_state.normalizer_params,
-                training_state.policy_params,
-                training_state.qr_params,
-                training_state.qc_params,
+                training_state.behavior_policy_params,
+                training_state.behavior_qr_params,
+                training_state.backup_qc_params,
                 training_state.model_params,
             )
             if store_buffer:
@@ -660,9 +668,9 @@ def training_epoch_with_timing(
     assert total_steps >= num_timesteps
     params = (
         training_state.normalizer_params,
-        training_state.policy_params,
-        training_state.qr_params,
-        training_state.qc_params,
+        training_state.behavior_policy_params,
+        training_state.behavior_qr_params,
+        training_state.backup_qc_params,
         training_state.model_params,
     )
     if store_buffer: