Add safety and terminations to the planning MDP

ss2r/algorithms/mbpo/model_env.py

@@ -4,12 +4,15 @@
 from brax.envs import base
 from brax.training.acme import running_statistics
 
+from ss2r.algorithms.sac.types import float32
+
 
 class ModelBasedEnv(envs.Env):
     """Environment wrapper that uses a learned model for predictions."""
 
     def __init__(
         self,
+        transitions,
         observation_size: int,
         action_size: int,
         model_network,
@@ -26,12 +29,36 @@ def __init__(
         self.safety_budget = safety_budget
         self._observation_size = observation_size
         self._action_size = action_size
+        self.transitions = transitions
 
     def reset(self, rng: jax.Array) -> base.State:
-        """Reset using the real environment."""
-        raise NotImplementedError(
-            "ModelBasedEnv does not support reset. Use a real environment for resetting."
+        sample_key, model_key = jax.random.split(rng)
+        indcs = jax.random.randint(
+            sample_key, (), 0, self.transitions.observation.shape[0]
+        )
+        transitions = float32(
+            jax.tree_util.tree_map(lambda x: x[indcs], self.transitions)
+        )
+        state = envs.State(
+            pipeline_state=None,
+            obs=transitions.observation,
+            reward=transitions.reward,
+            done=jnp.zeros_like(transitions.reward),
+            info={
+                "cumulative_cost": transitions.extras["state_extras"].get(
+                    "cumulative_cost", jnp.zeros_like(transitions.reward)
+                ),
+                "curr_discount": transitions.extras["state_extras"].get(
+                    "curr_discount", jnp.ones_like(transitions.reward)
+                ),
+                "truncation": jnp.zeros_like(transitions.reward),
+                "cost": transitions.extras["state_extras"].get(
+                    "cost", jnp.zeros_like(transitions.reward)
+                ),
+                "key": model_key,
+            },
         )
+        return state
 
     def step(self, state: base.State, action: jax.Array) -> base.State:
         """Step using the learned model."""
@@ -54,19 +81,30 @@ def step(self, state: base.State, action: jax.Array) -> base.State:
         state.info["truncation"] = truncation
         if "cumulative_cost" in state.info:
             prev_cumulative_cost = state.info["cumulative_cost"]
-            accumulated_cost_for_transition = prev_cumulative_cost + cost
-            if self.safety_budget < float("inf"):
-                done = jnp.where(
-                    accumulated_cost_for_transition > self.safety_budget,
-                    jnp.ones_like(done),
+            curr_discount = state.info.get("curr_discount", jnp.ones_like(reward))
+            accumulated_cost_for_transition = (
+                prev_cumulative_cost + curr_discount * cost
+            )
+            done = jnp.where(
+                accumulated_cost_for_transition > self.safety_budget,
+                jnp.ones_like(done),
+                done,
+            )
+
+            def reset_states(self, done, state, next_obs):
+                """Reset the state if done."""
+                key, reset_keys = jax.random.split(state.info["key"])
+                state.info["key"] = key
+                state.info["cumulative_cost"] = jnp.where(
                     done,
+                    jnp.zeros_like(reward),
+                    state.info["cumulative_cost"],
                 )
-            accumulated_cost_for_transition = jnp.where(
-                done > 0,
-                jnp.zeros_like(accumulated_cost_for_transition),
-                accumulated_cost_for_transition,
-            )
-            state.info["cumulative_cost"] = accumulated_cost_for_transition
+                next_obs = jnp.where(done, self.reset(reset_keys).obs, next_obs)
+                return state, next_obs
+
+            state, next_obs = reset_states(self, done, state, next_obs)
+
         state = state.replace(
             obs=next_obs,
             reward=reward,
@@ -89,6 +127,7 @@ def backend(self) -> str:
 
 
 def create_model_env(
+    transitions,
     model_network,
     model_params,
     observation_size: int,
@@ -99,6 +138,7 @@ def create_model_env(
 ) -> ModelBasedEnv:
     """Factory function to create a model-based environment."""
     return ModelBasedEnv(
+        transitions=transitions,
         model_network=model_network,
         observation_size=observation_size,
         action_size=action_size,

ss2r/algorithms/mbpo/train.py

@@ -37,6 +37,7 @@
 from ss2r.algorithms.mbpo.model_env import create_model_env
 from ss2r.algorithms.mbpo.training_step import make_training_step
 from ss2r.algorithms.mbpo.types import TrainingState
+from ss2r.algorithms.ppo.wrappers import TrackOnlineCosts
 from ss2r.algorithms.sac import gradients
 from ss2r.algorithms.sac.data import collect_single_step
 from ss2r.algorithms.sac.q_transforms import QTransformation, SACBase
@@ -190,6 +191,8 @@ def train(
     env = environment
     if wrap_env_fn is not None:
         env = wrap_env_fn(env)
+    if safe:
+        env = TrackOnlineCosts(env, safety_discounting)
     rng = jax.random.PRNGKey(seed)
     obs_size = env.observation_size
     action_size = env.action_size
@@ -226,6 +229,8 @@ def train(
     }
     if safe:
         extras["state_extras"]["cost"] = jnp.zeros(())  # type: ignore
+        extras["state_extras"]["cumulative_cost"] = jnp.zeros(())  # type: ignore
+        extras["state_extras"]["curr_discount"] = jnp.ones(())  # type: ignore
     dummy_transition = Transition(  # pytype: disable=wrong-arg-types  # jax-ndarray
         observation=dummy_obs,
         action=dummy_action,
@@ -314,7 +319,11 @@ def train(
     )
     extra_fields = ("truncation",)
     if safe:
-        extra_fields += ("cost",)  # type: ignore
+        extra_fields += (
+            "cost",
+            "cumulative_cost",
+            "curr_discount",
+        )  # type: ignore
 
     make_model_env = functools.partial(
         create_model_env,

ss2r/algorithms/mbpo/training_step.py

@@ -163,35 +163,15 @@ def run_experience_step(
     def generate_model_data(
         planning_env: ModelBasedEnv,
         policy: Policy,
-        transitions: Transition,
         sac_replay_buffer_state: ReplayBufferState,
         key: PRNGKey,
     ) -> ReplayBufferState:
         key_generate_unroll, cost_key, model_key, key_perm = jax.random.split(key, 4)
-        assert (
-            num_model_rollouts
-            <= transitions.observation.shape[0] * transitions.observation.shape[1]
-        ), "num_model_rollouts must be less than or equal to the number of transitions"
-        transitions = jax.tree_map(lambda x: x[:num_model_rollouts], transitions)
-        transitions = float32(transitions)
-        # FIXME: not zeros, use what happened in the real system
-        cumulative_cost = jax.random.uniform(
-            cost_key, (transitions.reward.shape[0],), minval=0.0, maxval=0.0
-        )
-        state = envs.State(
-            pipeline_state=None,
-            obs=transitions.observation,
-            reward=transitions.reward,
-            done=jnp.zeros_like(transitions.reward),
-            info={
-                "cumulative_cost": cumulative_cost,  # type: ignore
-                "truncation": jnp.zeros_like(cumulative_cost),
-                "cost": transitions.extras["state_extras"].get(
-                    "cost", jnp.zeros_like(cumulative_cost)
-                ),
-                "key": jnp.tile(model_key[None], (transitions.observation.shape[0], 1)),
-            },
-        )
+        keys = jax.random.split(key, num_model_rollouts + 2)
+        key = keys[0]
+        key_generate_unroll = keys[1]
+        rollout_keys = keys[2:]
+        state = planning_env.reset(rollout_keys)
         _, transitions = acting.generate_unroll(
             planning_env,
             state,
@@ -244,6 +224,9 @@ def training_step(
             training_key,
         ) = run_experience_step(training_state, env_state, model_buffer_state, key)
         model_buffer_state, transitions = model_replay_buffer.sample(model_buffer_state)
+        assert (
+            num_model_rollouts <= transitions.observation.shape[0]
+        ), "num_model_rollouts must be less than or equal to the number of transitions"
         # Change the front dimension of transitions so 'update_step' is called
         # grad_updates_per_step times by the scan.
         tmp_transitions = jax.tree_util.tree_map(
@@ -256,14 +239,15 @@ def training_step(
         planning_env = make_model_env(
             model_params=training_state.model_params,
             normalizer_params=training_state.normalizer_params,
+            transitions=transitions,
         )
         planning_env = VmapWrapper(planning_env)
         policy = make_policy(
             (training_state.normalizer_params, training_state.policy_params)
         )
         # Rollout trajectories from the sampled transitions
         sac_buffer_state = generate_model_data(
-            planning_env, policy, transitions, sac_buffer_state, training_key
+            planning_env, policy, sac_buffer_state, training_key
         )
         # Train SAC with model data
         sac_buffer_state, model_transitions = sac_replay_buffer.sample(sac_buffer_state)

ss2r/algorithms/ppo/wrappers.py

@@ -4,11 +4,16 @@
 
 
 class TrackOnlineCosts(Wrapper):
+    def __init__(self, env, cost_discount=1.0):
+        super().__init__(env)
+        self.cost_discount = cost_discount
+
     def reset(self, rng: jax.Array) -> State:
         reset_state = self.env.reset(rng)
         reset_state.info["cumulative_cost"] = reset_state.info.get(
             "cost", jnp.zeros_like(reset_state.reward)
         )
+        reset_state.info["curr_discount"] = jnp.ones_like(reset_state.reward)
         return reset_state
 
     def step(self, state: State, action: jax.Array) -> State:
@@ -19,7 +24,9 @@ def step(self, state: State, action: jax.Array) -> State:
         )
         nstate = self.env.step(state, action)
         cost = nstate.info.get("cost", jnp.zeros_like(nstate.reward))
-        nstate.info.update(cumulative_cost=cumulative_cost + cost)
+        curr_discount = nstate.info.get("curr_discount", jnp.ones_like(nstate.reward))
+        nstate.info.update(cumulative_cost=cumulative_cost + curr_discount * cost)
+        nstate.info.update(curr_discount=curr_discount * self.cost_discount)
         return nstate
 
 

ss2r/configs/experiment/cartpole_mbpo.yaml

@@ -6,7 +6,8 @@ defaults:
 
 training:
   num_timesteps: 150000
-  safe: false
+  safe: true
+  safety_budget: 100
   num_envs: 10
   train_domain_randomization: false
   eval_domain_randomization: false