Updates from deployable-rl (#19)

Author: yardenas

safe_opax/configs/agent/la_mbda.yaml

@@ -8,8 +8,10 @@ replay_buffer:
   capacity: 1000
 sentiment:
   ensemble_size: 5
-  model_initialization_scale: 1.
+  model_initialization_scale: 0.5
   critics_initialization_scale: 0.167
+  constraint_pessimism: null
+  objective_optimism: null
 model:
   hidden_size: 200
   stochastic_size: 60

safe_opax/configs/agent/penalizer/lbsgd.yaml

@@ -3,7 +3,8 @@ m_0: 1.2e4
 m_1: 1.2e4
 eta: 0.1
 eta_rate: 8e-6
-
+backup_lr: 1e-2
+  
 
 
 

safe_opax/configs/experiment/safe_sparse_cartpole.yaml

@@ -16,4 +16,7 @@ agent:
   exploration_strategy: opax
   exploration_steps: 1000000
   actor:
-    init_stddev: 0.001
+    init_stddev: 0.025
+  sentiment:
+    objective_optimism: 1.0
+    constraint_pessimism: 1.0

safe_opax/configs/experiment/safety_gym.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+defaults:
+  - override /environment: safe_adaptation_gym
+
+training:
+  epochs: 100
+  safe: true
+  action_repeat: 2
+
+agent:
+  exploration_steps: 0

safe_opax/la_mbda/exploration.py

@@ -44,7 +44,7 @@ def __init__(
             config.agent.model.stochastic_size + config.agent.model.deterministic_size,
             action_dim,
             key,
-            sentiment=identity,
+            objective_sentiment=identity,
         )
         self.reward_scale = config.agent.exploration_reward_scale
 

safe_opax/la_mbda/la_mbda.py

@@ -12,10 +12,11 @@
 from safe_opax.la_mbda.exploration import make_exploration
 from safe_opax.la_mbda.make_actor_critic import make_actor_critic
 from safe_opax.la_mbda.replay_buffer import ReplayBuffer
+from safe_opax.la_mbda.sentiment import Sentiment, UpperConfidenceBound, bayes
 from safe_opax.la_mbda.world_model import WorldModel, evaluate_model, variational_step
 from safe_opax.rl.epoch_summary import EpochSummary
 from safe_opax.rl.metrics import MetricsMonitor
-from safe_opax.rl.trajectory import TrajectoryData
+from safe_opax.rl.trajectory import TrajectoryData, Transition
 from safe_opax.rl.types import FloatArray, Report
 from safe_opax.rl.utils import Count, PRNGSequence, Until, add_to_buffer
 
@@ -51,6 +52,15 @@ def init(cls, batch_size: int, cell: rssm.RSSM, action_dim: int) -> "AgentState"
         return self
 
 
+def make_sentiment(alpha) -> Sentiment:
+    if alpha is None or alpha == 0.0:
+        return bayes
+    elif alpha > 0.0:
+        return UpperConfidenceBound(alpha)
+    else:
+        raise ValueError(f"Invalid alpha: {alpha}")
+
+
 class LaMBDA:
     def __init__(
         self,
@@ -86,6 +96,8 @@ def __init__(
             config.agent.model.stochastic_size + config.agent.model.deterministic_size,
             action_dim,
             next(self.prng),
+            make_sentiment(self.config.agent.sentiment.objective_optimism),
+            make_sentiment(self.config.agent.sentiment.constraint_pessimism),
         )
         self.exploration = make_exploration(
             config,
@@ -137,6 +149,9 @@ def observe(self, trajectory: TrajectoryData) -> None:
         )
         self.state = jax.tree_map(lambda x: jnp.zeros_like(x), self.state)
 
+    def observe_transition(self, transition: Transition) -> None:
+        pass
+
     def update(self):
         total_steps = self.config.agent.update_steps
         for batch in self.replay_buffer.sample(total_steps):

safe_opax/la_mbda/lbsgd.py

@@ -27,7 +27,7 @@ def compute_lr(alpha_1, g, grad_f_1, m_0, m_1, eta):
     m_2 = (
         m_0
         + 10.0 * eta * (m_1 / (alpha_1 + _EPS))
-        + 8.0 * eta * (theta_1 / alpha_1 + _EPS) ** 2
+        + 8.0 * eta * (theta_1 / (alpha_1 + _EPS)) ** 2
     )
     rhs = 1.0 / m_2
     return jnp.minimum(lhs, rhs), (lhs, rhs)

safe_opax/la_mbda/make_actor_critic.py

@@ -10,7 +10,15 @@
 _LOG = logging.getLogger(__name__)
 
 
-def make_actor_critic(cfg, safe, state_dim, action_dim, key, sentiment=bayes):
+def make_actor_critic(
+    cfg,
+    safe,
+    state_dim,
+    action_dim,
+    key,
+    objective_sentiment=bayes,
+    constraint_sentiment=bayes,
+):
     # Account for the the discount factor in the budget.
     episode_safety_budget = (
         (
@@ -29,6 +37,7 @@ def make_actor_critic(cfg, safe, state_dim, action_dim, key, sentiment=bayes):
                 cfg.agent.penalizer.eta,
                 cfg.agent.penalizer.eta_rate,
                 cfg.agent.actor_optimizer.lr,
+                cfg.agent.penalizer.backup_lr,
             )
         elif cfg.agent.penalizer.name == "lagrangian":
             penalizer = AugmentedLagrangianPenalizer(
@@ -56,5 +65,6 @@ def make_actor_critic(cfg, safe, state_dim, action_dim, key, sentiment=bayes):
         safety_budget=episode_safety_budget,
         penalizer=penalizer,
         key=key,
-        objective_sentiment=sentiment,
+        objective_sentiment=objective_sentiment,
+        constraint_sentiment=constraint_sentiment,
     )

safe_opax/la_mbda/safe_actor_critic.py

@@ -26,6 +26,8 @@ class ActorEvaluation(NamedTuple):
     constraint: jax.Array
     safe: jax.Array
     priors: ShiftScale
+    reward_stddev: jax.Array
+    cost_stddev: jax.Array
 
 
 class Penalizer(Protocol):
@@ -59,6 +61,7 @@ def __init__(
         key: jax.Array,
         penalizer: Penalizer,
         objective_sentiment: Sentiment,
+        constraint_sentiment: Sentiment,
     ):
         actor_key, critic_key, safety_critic_key = jax.random.split(key, 3)
         self.actor = ContinuousActor(
@@ -84,18 +87,18 @@ def __init__(
         self.lambda_ = lambda_
         self.safety_discount = safety_discount
         self.safety_budget = safety_budget
-        self.update_fn = batched_update_safe_actor_critic
         self.penalizer = penalizer
         self.objective_sentiment = objective_sentiment
+        self.constraint_sentiment = constraint_sentiment
 
     def update(
         self,
         model: Model,
         initial_states: jax.Array,
         key: jax.Array,
     ) -> dict[str, float]:
-        actor_critic_fn = partial(self.update_fn, model.sample)
-        results: SafeActorCriticStepResults = actor_critic_fn(
+        results: SafeActorCriticStepResults = update_safe_actor_critic(
+            model.sample,
             self.horizon,
             initial_states,
             self.actor,
@@ -115,6 +118,7 @@ def update(
             self.penalizer,
             self.penalizer.state,
             self.objective_sentiment,
+            self.constraint_sentiment,
         )
         self.actor = results.new_actor
         self.critic = results.new_critic
@@ -196,6 +200,7 @@ def evaluate_actor(
     lambda_: float,
     safety_budget: float,
     objective_sentiment: Sentiment,
+    constraint_sentiment: Sentiment,
 ) -> ActorEvaluation:
     trajectories, priors = rollout_fn(horizon, initial_states, key, actor.act)
     next_step = lambda x: x[:, 1:]
@@ -207,9 +212,7 @@ def evaluate_actor(
         bootstrap_values, rewards, discount, lambda_
     )
     bootstrap_safety_values = nest_vmap(safety_critic, 2, eqx.filter_vmap)(next_states)
-    # TODO (yarden): make costs use their own sentiments when working
-    # on safety.
-    costs = current_step(trajectories.cost.mean(1))
+    costs = current_step(constraint_sentiment(trajectories.cost))
     safety_lambda_values = eqx.filter_vmap(compute_lambda_values)(
         bootstrap_safety_values,
         costs,
@@ -228,9 +231,16 @@ def evaluate_actor(
         constraint,
         jnp.greater(constraint, 0.0),
         priors,
+        rewards.std(1).mean(),
+        costs.std(1).mean(),
     )
 
 
+@eqx.filter_jit
+@apply_mixed_precision(
+    target_module_names=["critic", "safety_critic", "actor", "rollout_fn"],
+    target_input_names=["initial_states"],
+)
 def update_safe_actor_critic(
     rollout_fn: RolloutFn,
     horizon: int,
@@ -252,13 +262,15 @@ def update_safe_actor_critic(
     penalty_fn: Penalizer,
     penalty_state: Any,
     objective_sentiment: Sentiment,
+    constraint_sentiment: Sentiment,
 ) -> SafeActorCriticStepResults:
+    vmapped_rollout_fn = jax.vmap(rollout_fn, (None, 0, None, None))
     actor_grads, new_penalty_state, evaluation, metrics = penalty_fn(
         lambda actor: evaluate_actor(
             actor,
             critic,
             safety_critic,
-            rollout_fn,
+            vmapped_rollout_fn,
             horizon,
             initial_states,
             key,
@@ -267,6 +279,7 @@ def update_safe_actor_critic(
             lambda_,
             safety_budget,
             objective_sentiment,
+            constraint_sentiment,
         ),
         penalty_state,
         actor,
@@ -292,9 +305,11 @@ def update_safe_actor_critic(
     new_safety_critic, new_safety_critic_state = safety_critic_learner.grad_step(
         safety_critic, grads, safety_critic_learning_state
     )
-    metrics["agent/epistemic_uncertainty"] = normalized_epistemic_uncertainty(
+    metrics["agent/sentiment/epistemic_uncertainty"] = normalized_epistemic_uncertainty(
         evaluation.priors, 1
     ).mean()
+    metrics["agent/sentiment/reward_stddev"] = evaluation.reward_stddev
+    metrics["agent/sentiment/cost_stddev"] = evaluation.cost_stddev
     return SafeActorCriticStepResults(
         new_actor,
         new_critic,
@@ -313,58 +328,6 @@ def update_safe_actor_critic(
     )
 
 
-@eqx.filter_jit
-@apply_mixed_precision(
-    target_module_names=["critic", "safety_critic", "actor", "rollout_fn"],
-    target_input_names=["initial_states"],
-)
-def batched_update_safe_actor_critic(
-    rollout_fn: RolloutFn,
-    horizon: int,
-    initial_states: jax.Array,
-    actor: ContinuousActor,
-    critic: Critic,
-    safety_critic: Critic,
-    actor_learning_state: OptState,
-    critic_learning_state: OptState,
-    safety_critic_learning_state: OptState,
-    actor_learner: Learner,
-    critic_learner: Learner,
-    safety_critic_learner: Learner,
-    key: jax.Array,
-    discount: float,
-    safety_discount: float,
-    lambda_: float,
-    safety_budget: float,
-    penalty_fn: Penalizer,
-    penalty_state: Any,
-    objective_sentiment: Sentiment,
-) -> SafeActorCriticStepResults:
-    vmapped_rollout_fn = jax.vmap(rollout_fn, (None, 0, None, None))
-    return update_safe_actor_critic(
-        vmapped_rollout_fn,
-        horizon,
-        initial_states,
-        actor,
-        critic,
-        safety_critic,
-        actor_learning_state,
-        critic_learning_state,
-        safety_critic_learning_state,
-        actor_learner,
-        critic_learner,
-        safety_critic_learner,
-        key,
-        discount,
-        safety_discount,
-        lambda_,
-        safety_budget,
-        penalty_fn,
-        penalty_state,
-        objective_sentiment,
-    )
-
-
 def compute_discount(factor, length):
     d = jnp.cumprod(factor * jnp.ones((length - 1,)))
     d = jnp.concatenate([jnp.ones((1,)), d])

safe_opax/la_mbda/sentiment.py

@@ -16,6 +16,23 @@ def bayes(values: jax.Array) -> jax.Array:
     return values.mean(1)
 
 
+class UpperConfidenceBound(Sentiment):
+    def __init__(self, alpha: float = 1.0):
+        self.alpha = alpha
+
+    def __call__(self, values: jax.Array) -> jax.Array:
+        return upper_confidence_bound(values, self.alpha)
+
+
+def upper_confidence_bound(
+    values: jax.Array, alpha: float, stop_gradient: bool = True
+) -> jax.Array:
+    stddev = jnp.std(values, axis=1)
+    if stop_gradient:
+        stddev = jax.lax.stop_gradient(stddev)
+    return jnp.mean(values, axis=1) + alpha * stddev
+
+
 def _emprirical_estimate(
     values: jax.Array, reduce_fn: Callable[[jax.Array], jax.Array]
 ) -> jax.Array: