Add a constraint cartpole environment (#14)

* Add a constraint cartpole environment

* Setup experiment configuration

* Cost rate -> cost return

* Bug fixes with lbsgd

* Scale optimism rewards

* Fix bug in rewards scale

* Fix baking state of exploration into policy bug

* Zero out bias

* Initially safe policy

* Reset metrics completely

* Log lhs rhs

* Fix bugs in rhs computation, tune lrs

* Log safety budget

* Fix learning rate scaling in lbsgd

Author: yardenas

poetry.lock

@@ -203,23 +203,44 @@ def __init__(self, env: Env, cost_multiplier: float = 0):
         self.cost_multiplier = cost_multiplier
 
     def step(self, action):
-        action_cost = self.cost_multiplier * (1 - tolerance(action, (-0.1, 0.1), 0.1))[0]
+        action_cost = (
+            self.cost_multiplier * (1 - tolerance(action, (-0.1, 0.1), 0.1))[0]
+        )
         observation, reward, terminal, truncated, info = self.env.step(action)
         return observation, reward - action_cost, terminal, truncated, info
 
     def __getattr__(self, name):
         return getattr(self.env, name)
 
 
+class ConstraintWrapper:
+    def __init__(self, env: Env, slider_position_bound: float):
+        self.env = env
+        self.physics = env.env.physics
+        self.slider_position_bound = slider_position_bound
+
+    def step(self, action):
+        observation, reward, terminal, truncated, info = self.env.step(action)
+        slider_pos = self.physics.cart_position().copy()
+        cost = float(np.abs(slider_pos) >= self.slider_position_bound)
+        info["cost"] = cost
+        return observation, reward, terminal, truncated, info
+
+    def __getattr__(self, name):
+        return getattr(self.env, name)
+
+
 def make(cfg: DictConfig) -> EnvironmentFactory:
     def make_env():
         domain_name, task_cfg = get_domain_and_task(cfg)
-        if task_cfg.task == "swingup_sparse_hard":
+        if task_cfg.task in ["swingup_sparse_hard", "safe_swingup_sparse_hard"]:
             task = "swingup_sparse"
         else:
             task = task_cfg.task
         env = DMCWrapper(domain_name, task)
-        if task_cfg.task == "swingup_sparse_hard":
+        if "safe" in task_cfg.task:
+            env = ConstraintWrapper(env, task_cfg.slider_position_bound)
+        if task_cfg.task in ["swingup_sparse_hard", "safe_swingup_sparse_hard"]:
             env = ActionCostWrapper(env, cost_multiplier=task_cfg.cost_multiplier)
         if task_cfg.image_observation.enabled:
             env = ImageObservation(
@@ -245,6 +266,7 @@ def make_env():
     ("dm_cartpole", "swingup"),
     ("dm_cartpole", "swingup_sparse"),
     ("dm_cartpole", "swingup_sparse_hard"),
+    ("dm_cartpole", "safe_swingup_sparse_hard"),
     ("dm_humanoid", "stand"),
     ("dm_humanoid", "walk"),
     ("dm_manipulator", "bring_ball"),

safe_opax/benchmark_suites/dm_control/__init__.py

@@ -1,5 +1,5 @@
 defaults:
-  - penalizer/lagrangian
+  - penalizer: lagrangian
 
 name: lambda
 replay_buffer:
@@ -18,6 +18,7 @@ actor:
   n_layers: 4
   hidden_size: 400
   init_stddev: 5.
+  initialization_scale: 0.01
 critic:
   n_layers: 3
   hidden_size: 400
@@ -49,4 +50,5 @@ kl_mix: 0.8
 safety_slack: 0.
 evaluate_model: false
 exploration_strategy: uniform
-exploration_steps: 5000
+exploration_steps: 5000
+exploration_reward_scale: 10.0

safe_opax/configs/agent/la_mbda.yaml

@@ -1,6 +1,6 @@
 name: lbsgd
-m_0: 5e6
-m_1: 5e4
+m_0: 1.2e4
+m_1: 1.2e4
 eta: 0.1
 eta_rate: 8e-6
 

safe_opax/configs/agent/penalizer/lbsgd.yaml

@@ -6,3 +6,4 @@ dm_cartpole:
     image_format: "channels_first"
     visualize_reward: true
   cost_multiplier: 0.2
+  slider_position_bound: 0.5

safe_opax/configs/environment/dm_cartpole.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+defaults:
+  - override /environment: dm_cartpole
+  - override /agent/penalizer: lbsgd
+
+environment:
+  dm_cartpole:
+    task: safe_swingup_sparse_hard
+
+training:
+  epochs: 100
+  safe: true
+  action_repeat: 2
+  safety_budget: 100
+
+agent:
+  exploration_strategy: opax
+  exploration_steps: 1000000
+  actor:
+    init_stddev: 0.001

safe_opax/configs/experiment/safe_sparse_cartpole.yaml

@@ -20,6 +20,7 @@ def __init__(
         action_dim: int,
         hidden_size: int,
         init_stddev: float,
+        initialization_scale: float,
         *,
         key: jax.Array,
     ):
@@ -31,7 +32,8 @@ def __init__(
                 n_layers + 1,
                 key=key,
                 activation=jnn.elu,
-            )
+            ),
+            weight_scale=initialization_scale,
         )
         self.init_stddev = init_stddev
 

safe_opax/la_mbda/actor_critic.py

@@ -4,13 +4,10 @@
 from safe_opax.la_mbda.opax_bridge import OpaxBridge
 from safe_opax.la_mbda.make_actor_critic import make_actor_critic
 from safe_opax.la_mbda.sentiment import identity
-from safe_opax.rl.types import Model
+from safe_opax.rl.types import Model, Policy
 
 
 class Exploration:
-    def __call__(self, state: jax.Array, key: jax.Array) -> jax.Array:
-        raise NotImplementedError("Must be implemented by subclass")
-
     def update(
         self,
         model: Model,
@@ -19,6 +16,9 @@ def update(
     ) -> dict[str, float]:
         return {}
 
+    def get_policy(self) -> Policy:
+        raise NotImplementedError("Must be implemented by subclass")
+
 
 def make_exploration(
     config: DictConfig, action_dim: int, key: jax.Array
@@ -46,31 +46,32 @@ def __init__(
             key,
             sentiment=identity,
         )
+        self.reward_scale = config.agent.exploration_reward_scale
 
     def update(
         self,
         model: Model,
         initial_states: jax.Array,
         key: jax.Array,
     ) -> dict[str, float]:
-        model = OpaxBridge(model)
+        model = OpaxBridge(model, self.reward_scale)
         outs = self.actor_critic.update(model, initial_states, key)
+        outs = {f"{_append_opax(k)}": v for k, v in outs.items()}
+        return outs
 
-        def append_opax(string):
-            parts = string.split("/")
-            parts.insert(2, "opax")
-            return "/".join(parts)
+    def get_policy(self) -> Policy:
+        return self.actor_critic.actor.act
 
-        outs = {f"{append_opax(k)}": v for k, v in outs.items()}
-        return outs
 
-    def __call__(self, state: jax.Array, key: jax.Array) -> jax.Array:
-        return self.actor_critic.actor.act(state, key)
+def _append_opax(string):
+    parts = string.split("/")
+    parts.insert(2, "opax")
+    return "/".join(parts)
 
 
 class UniformExploration(Exploration):
     def __init__(self, action_dim: int):
         self.action_dim = action_dim
 
-    def __call__(self, state: jax.Array, key: jax.Array) -> jax.Array:
-        return jax.random.uniform(key, (self.action_dim,))
+    def get_policy(self) -> Policy:
+        return lambda state, key: jax.random.uniform(key, (self.action_dim,))

safe_opax/la_mbda/exploration.py

@@ -115,7 +115,9 @@ def __call__(
         if train and self.should_train() and not self.replay_buffer.empty:
             self.update()
         policy_fn = (
-            self.exploration if self.should_explore() else self.actor_critic.actor.act
+            self.exploration.get_policy()
+            if self.should_explore()
+            else self.actor_critic.actor.act
         )
         self.should_explore.tick()
         actions, self.state = policy(

safe_opax/la_mbda/la_mbda.py

@@ -7,57 +7,56 @@
 import jax.numpy as jnp
 from jaxtyping import PyTree
 
+from safe_opax.common.mixed_precision import apply_dtype
 from safe_opax.common.pytree_utils import pytrees_unstack
 from safe_opax.la_mbda.actor_critic import ContinuousActor
 from safe_opax.la_mbda.safe_actor_critic import ActorEvaluation
 
+_EPS = 1e-8
+
 
 class LBSGDState(NamedTuple):
     eta: jax.Array
 
 
-def compute_lr(constraint, loss_grads, constraint_grads, m_0, m_1, eta):
-    constraint_grads, _ = jax.flatten_util.ravel_pytree(constraint_grads)  # type: ignore
-    loss_grads, _ = jax.flatten_util.ravel_pytree(loss_grads)  # type: ignore
-    projection = constraint_grads.dot(loss_grads)
-    lhs = (
-        constraint
-        / (
-            2.0 * jnp.abs(projection) / jnp.linalg.norm(loss_grads)
-            + jnp.sqrt(constraint * m_1 + 1e-8)
-        )
-        / (jnp.linalg.norm(loss_grads) + 1e-8)
-    )
+def compute_lr(alpha_1, g, grad_f_1, m_0, m_1, eta):
+    grad_f_1, _ = jax.flatten_util.ravel_pytree(grad_f_1)
+    g, _ = jax.flatten_util.ravel_pytree(g)
+    theta_1 = grad_f_1.dot(g / (jnp.linalg.norm(g) + _EPS))
+    lhs = alpha_1 / (2.0 * jnp.abs(theta_1) + jnp.sqrt(alpha_1 * m_1 + _EPS))
     m_2 = (
         m_0
-        + 10.0 * eta * (m_1 / (constraint + 1e-8))
-        + 8.0
-        * eta
-        * jnp.linalg.norm(projection) ** 2
-        / ((jnp.linalg.norm(loss_grads) * constraint) ** 2)
+        + 10.0 * eta * (m_1 / (alpha_1 + _EPS))
+        + 8.0 * eta * (theta_1 / alpha_1 + _EPS) ** 2
     )
     rhs = 1.0 / m_2
-    return jnp.minimum(lhs, rhs)
+    return jnp.minimum(lhs, rhs), (lhs, rhs)
 
 
 def lbsgd_update(
-    state: LBSGDState, updates: PyTree, eta_rate: float, m_0: float, m_1: float
-) -> tuple[PyTree, LBSGDState]:
+    state: LBSGDState,
+    updates: PyTree,
+    eta_rate: float,
+    m_0: float,
+    m_1: float,
+    base_lr: float,
+    backup_lr: float,
+) -> tuple[PyTree, LBSGDState, tuple[float, ...]]:
     def happy_case():
-        lr = compute_lr(constraint, loss_grads, constraints_grads, m_0, m_1, eta_t)
+        lr, (lhs, rhs) = compute_lr(alpha_1, g, grad_f_1, m_0, m_1, eta_t)
         new_eta = eta_t / eta_rate
-        updates = jax.tree_map(lambda x: x * lr, loss_grads)
-        return updates, LBSGDState(new_eta)
+        updates = jax.tree_map(lambda x: x * lr / base_lr, g)
+        return updates, LBSGDState(new_eta), (lr, lhs, rhs)
 
     def fallback():
         # Taking the negative gradient of the constraints to minimize the costs
-        updates = jax.tree_map(lambda x: x * -1.0, constraints_grads)
-        return updates, LBSGDState(eta_t)
+        updates = jax.tree_map(lambda x: x * backup_lr, grad_f_1)
+        return updates, LBSGDState(eta_t), (0.0, 0.0, 0.0)
 
-    loss_grads, constraints_grads, constraint = updates
+    g, grad_f_1, alpha_1 = updates
     eta_t = state.eta
     return jax.lax.cond(
-        jnp.greater(constraint, 0.0),
+        jnp.greater(alpha_1, _EPS),
         happy_case,
         fallback,
     )
@@ -66,17 +65,27 @@ def fallback():
 def jacrev(f, has_aux=False):
     def jacfn(x):
         y, vjp_fn, aux = eqx.filter_vjp(f, x, has_aux=has_aux)  # type: ignore
-        (J,) = eqx.filter_vmap(vjp_fn, in_axes=0)(jnp.eye(len(y)))
+        (J,) = eqx.filter_vmap(vjp_fn, in_axes=eqx.if_array(0))(jnp.eye(len(y)))
         return J, aux
 
     return jacfn
 
 
 class LBSGDPenalizer:
-    def __init__(self, m_0, m_1, eta, eta_rate) -> None:
+    def __init__(
+        self,
+        m_0: float,
+        m_1: float,
+        eta: float,
+        eta_rate: float,
+        base_lr: float,
+        backup_lr: float = 1e-2,
+    ) -> None:
         self.m_0 = m_0
         self.m_1 = m_1
         self.eta_rate = eta_rate + 1.0
+        self.base_lr = base_lr
+        self.backup_lr = backup_lr
         self.state = LBSGDState(eta)
 
     def __call__(
@@ -87,19 +96,26 @@ def __call__(
     ) -> tuple[PyTree, Any, ActorEvaluation, dict[str, jax.Array]]:
         def evaluate_helper(actor):
             evaluation = evaluate(actor)
-            outs = jnp.stack([evaluation.loss, evaluation.constraint])
+            loss = evaluation.loss - state.eta * jnp.log(evaluation.constraint)
+            outs = jnp.stack([loss, -evaluation.constraint])
             return outs, evaluation
 
         jacobian, rest = jacrev(evaluate_helper, has_aux=True)(actor)
-        loss_grads, constraint_grads = pytrees_unstack(jacobian)
-        updates, state = lbsgd_update(
+        g, grad_f_1 = pytrees_unstack(jacobian)
+        alpha = rest.constraint
+        updates, state, (lr, lhs, rhs) = lbsgd_update(
             state,
-            (loss_grads, constraint_grads, rest.constraint),
+            apply_dtype((g, grad_f_1, alpha), jnp.float32),
             self.eta_rate,
             self.m_0,
             self.m_1,
+            self.base_lr,
+            self.backup_lr,
         )
         metrics = {
-            "agent/lbsgd/eta": state.eta,
+            "agent/lbsgd/eta": jnp.asarray(state.eta),
+            "agent/lbsgd/lr": jnp.asarray(lr),
+            "agent/lbsgd/lhs": jnp.asarray(lhs),
+            "agent/lbsgd/rhs": jnp.asarray(rhs),
         }
         return updates, state, rest, metrics