Constraint evaluation (#11)

* Cleanup

* Add evaluator class and some cleanups

* Add constraint evaluator to sac

* Log all metrics

* Fix some bugs in constraints

Author: yardenas

ss2r/algorithms/sac/train.py

@@ -35,6 +35,7 @@
 
 import ss2r.algorithms.sac.losses as sac_losses
 import ss2r.algorithms.sac.networks as sac_networks
+from ss2r.rl.evaluation import ConstraintsEvaluator
 
 Metrics: TypeAlias = types.Metrics
 Transition: TypeAlias = types.Transition
@@ -534,7 +535,7 @@ def training_epoch_with_timing(
             randomization_fn=vf_randomization_fn,
         )
 
-    evaluator = acting.Evaluator(
+    evaluator = ConstraintsEvaluator(
         eval_env,
         functools.partial(make_policy, deterministic=deterministic_eval),
         num_eval_envs=num_eval_envs,

ss2r/benchmark_suites/rccar/rccar.py

@@ -159,7 +159,7 @@ def __init__(
         margin_factor: float = 10.0,
         max_throttle: float = 1.0,
         dt: float = 1 / 30.0,
-        obstacle: tuple[float, float, float] = (-0.75, -0.75, 0.2),
+        obstacle: tuple[float, float, float] = (0.75, -0.75, 0.2),
     ):
         """
         Race car simulator environment
@@ -173,21 +173,21 @@ def __init__(
             car_model_params: dictionary of car model parameters that overwrite the default values
             seed: random number generator seed
         """
-        self._goal = jnp.array([0.0, 0.0, 0.0])
+        self.goal = jnp.array([0.0, 0.0, 0.0])
         self.obstacle = tuple(obstacle)
         self._init_pose = jnp.array([1.42, -1.04, jnp.pi])
-        self._angle_idx = 2
+        self.angle_idx = 2
         self._obs_noise_stds = OBS_NOISE_STD_SIM_CAR
         self.dim_action = (2,)
-        self._dt = dt
+        self.dt = dt
         self.dim_state = (7,) if encode_angle else (6,)
         self.encode_angle = encode_angle
         self.max_throttle = jnp.clip(max_throttle, 0.0, 1.0)
-        self.dynamics_model = RaceCarDynamics(dt=self._dt)
+        self.dynamics_model = RaceCarDynamics(dt=self.dt)
         self.sys = CarParams(**car_model_params)
         self.use_obs_noise = use_obs_noise
         self.reward_model = RCCarEnvReward(
-            goal=self._goal,
+            goal=self.goal,
             ctrl_cost_weight=ctrl_cost_weight,
             encode_angle=self.encode_angle,
             margin_factor=margin_factor,
@@ -205,7 +205,7 @@ def _obs(self, state: jnp.array, rng: jax.random.PRNGKey) -> jnp.array:
             obs = state
         # encode angle to sin(theta) ant cos(theta) if desired
         if self.encode_angle:
-            obs = encode_angles(obs, self._angle_idx)
+            obs = encode_angles(obs, self.angle_idx)
         assert (obs.shape[-1] == 7 and self.encode_angle) or (
             obs.shape[-1] == 6 and not self.encode_angle
         )
@@ -241,7 +241,7 @@ def step(self, state: State, action: jax.Array) -> State:
         action = action.at[0].set(self.max_throttle * action[0])
         obs = state.obs
         if self.encode_angle:
-            dynamics_state = decode_angles(obs, self._angle_idx)
+            dynamics_state = decode_angles(obs, self.angle_idx)
         next_dynamics_state = self.dynamics_model.step(dynamics_state, action, self.sys)
         # FIXME (yarden): hard-coded key is bad here.
         next_obs = self._obs(next_dynamics_state, rng=jax.random.PRNGKey(0))
@@ -259,10 +259,6 @@ def step(self, state: State, action: jax.Array) -> State:
         )
         return next_state
 
-    @property
-    def dt(self):
-        return self._dt
-
     @property
     def observation_size(self) -> int:
         if self.encode_angle:
@@ -289,8 +285,8 @@ def render(env, policy, steps, rng):
     trajectory = jax.tree_map(lambda x: x[:, 0], trajectory.obs)
     if env.encode_angle:
         trajectory = decode_angles(trajectory, 2)
-
     obstacle_position, obstacle_radius = env.obstacle[:2], env.obstacle[2]
+    obstacle_position = jnp.array([obstacle_position[1], -obstacle_position[0]])
 
     def draw_scene(timestep):
         # Create a figure and axis

ss2r/common/learner.py

@@ -11,4 +11,4 @@ ctrl_cost_weight: 0.005
 margin_factor: 20.0
 max_throttle: 1.0
 use_obs_noise: false
-obstacle: [-0.75, -0.75, 0.2]  # x, y, radius
+obstacle: [0.75, -0.75, 0.2]  # x, y, radius

ss2r/rl/logging.py renamed to ss2r/common/logging.py

@@ -0,0 +1,78 @@
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+from brax.envs.base import Env, State
+from brax.envs.wrappers.training import EvalMetrics, EvalWrapper
+from brax.training.acting import Evaluator, generate_unroll
+from brax.training.types import Policy, PolicyParams, PRNGKey
+
+
+class ConstraintEvalWrapper(EvalWrapper):
+    def reset(self, rng: jax.Array) -> State:
+        reset_state = self.env.reset(rng)
+        reset_state.metrics["reward"] = reset_state.reward
+        reset_state.metrics["cost"] = reset_state.info.get("cost", jnp.array(0.0))
+        eval_metrics = EvalMetrics(
+            episode_metrics=jax.tree_util.tree_map(jnp.zeros_like, reset_state.metrics),
+            active_episodes=jnp.ones_like(reset_state.reward),
+            episode_steps=jnp.zeros_like(reset_state.reward),
+        )
+        reset_state.info["eval_metrics"] = eval_metrics
+        return reset_state
+
+    def step(self, state: State, action: jax.Array) -> State:
+        state_metrics = state.info["eval_metrics"]
+        if not isinstance(state_metrics, EvalMetrics):
+            raise ValueError(f"Incorrect type for state_metrics: {type(state_metrics)}")
+        del state.info["eval_metrics"]
+        nstate = self.env.step(state, action)
+        nstate.metrics["reward"] = nstate.reward
+        nstate.metrics["cost"] = nstate.info.get("cost", jnp.array(0.0))
+        episode_steps = jnp.where(
+            state_metrics.active_episodes,
+            nstate.info["steps"],
+            state_metrics.episode_steps,
+        )
+        episode_metrics = jax.tree_util.tree_map(
+            lambda a, b: a + b * state_metrics.active_episodes,
+            state_metrics.episode_metrics,
+            nstate.metrics,
+        )
+        active_episodes = state_metrics.active_episodes * (1 - nstate.done)
+        eval_metrics = EvalMetrics(
+            episode_metrics=episode_metrics,
+            active_episodes=active_episodes,
+            episode_steps=episode_steps,
+        )
+        nstate.info["eval_metrics"] = eval_metrics
+        return nstate
+
+
+class ConstraintsEvaluator(Evaluator):
+    def __init__(
+        self,
+        eval_env: Env,
+        eval_policy_fn: Callable[[PolicyParams], Policy],
+        num_eval_envs: int,
+        episode_length: int,
+        action_repeat: int,
+        key: jax.Array,
+    ):
+        self._key = key
+        self._eval_walltime = 0.0
+        eval_env = ConstraintEvalWrapper(eval_env)
+
+        def generate_eval_unroll(policy_params: PolicyParams, key: PRNGKey) -> State:
+            reset_keys = jax.random.split(key, num_eval_envs)
+            eval_first_state = eval_env.reset(reset_keys)
+            return generate_unroll(
+                eval_env,
+                eval_first_state,
+                eval_policy_fn(policy_params),
+                key,
+                unroll_length=episode_length // action_repeat,
+            )[0]
+
+        self._generate_eval_unroll = jax.jit(generate_eval_unroll)
+        self._steps_per_unroll = episode_length * num_eval_envs