fix: PPO extracted and integrated with jax

ruff.toml

@@ -347,7 +347,6 @@ extend-ignore = [
   # "PLR1705",  # no-else-return
   # "PLR1706",  # consider-using-ternary
   # "PLR1707",  # trailing-comma-tuple
-  "PLR1708",  # stop-iteration-return
   # "PLR1709",  # simplify-boolean-expression
   # "PLR1710",  # inconsistent-return-statements
   "PLR1711",  # useless-return

src/ppo.py

@@ -1,80 +1,48 @@
-# docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_atari_envpool_xla_jaxpy
+from functools import partial
 
+import flax
 import jax
 import jax.numpy as jnp
-from flax.training.train_state import TrainState
 
-@jax.jit
-def get_action_and_value2(
-    params: flax.core.FrozenDict,
-    x: np.ndarray,
-    action: np.ndarray,
-):
-    """calculate value, logprob of supplied `action`, and entropy"""
-    hidden = network.apply(params.network_params, x)
-    # assume that actor returns mean and log std over continuos action space, why log?, better for..?, research this
-    mean, log_std = actor.apply(params.actor_params, hidden)
-    std = jnp.exp(log_std)
 
-    # compute logprob of the given action
+# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
-    var = std ** 2
+# with our need to seperate concerns
-    logprob = -0.5 * (((action - mean) ** 2) / var + 2 * log_std + jnp.log(2 * jnp.pi))
+class PPO:
-    logprob = logprob.sum(-1)
+    def __init__(self, args, network, actor, critic):
+        self.args = args
+        self.network = network
+        self.actor = actor
+        self.critic = critic
 
-    # Validate that this is a good entropy (check other ppo implementations)
+        self.ppo_loss_grad_fn = jax.value_and_grad(
-    entropy = 0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std
+            partial(ppo_loss, args=args, network=network, actor=actor, critic=critic),
-    entropy = entropy.sum(-1)
+            has_aux=True,
+        )
 
-    value = critic.apply(params.critic_params, hidden).squeeze()
+    # This PPO class should be initialized only once,
-    return logprob, entropy, value
+    # or this function will need to recompile
+    @partial(jax.jit, static_argnums=0)
+    def update_ppo(self, agent_state, storage, key):
+        args = self.args
+        ppo_loss_grad_fn = self.ppo_loss_grad_fn
 
-def ppo_loss(params, x, a, logp, mb_advantages, mb_returns):
+        def update_epoch(carry, _):
-    newlogprob, entropy, newvalue = get_action_and_value2(params, x, a)
-    logratio = newlogprob - logp
-    ratio = jnp.exp(logratio)
-    approx_kl = ((ratio - 1) - logratio).mean()
-
-    if args.norm_adv:
-        mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
-
-    # Policy loss
-    pg_loss1 = -mb_advantages * ratio
-    pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
-    pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
-
-    # Value loss
-    v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
-
-    entropy_loss = entropy.mean()
-    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
-    return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
-
-ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True)
-
-@jax.jit
-def update_ppo(
-    agent_state: TrainState,
-    storage: Storage,
-    key: jax.random.PRNGKey,
-):
-    def update_epoch(carry, unused_inp):
             agent_state, key = carry
             key, subkey = jax.random.split(key)
 
             def flatten(x):
                 return x.reshape((-1,) + x.shape[2:])
 
-        # taken from: https://github.com/google/brax/blob/main/brax/training/agents/ppo/train.py
+            def convert_data(x):
-        def convert_data(x: jnp.ndarray):
                 x = jax.random.permutation(subkey, x)
-            x = jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
+                return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
-            return x
 
-        flatten_storage = jax.tree_map(flatten, storage)
+            flatten_storage = jax.tree.map(flatten, storage)
-        shuffled_storage = jax.tree_map(convert_data, flatten_storage)
+            shuffled_storage = jax.tree.map(convert_data, flatten_storage)
 
             def update_minibatch(agent_state, minibatch):
-            (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
+                (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = (
+                    ppo_loss_grad_fn(
                         agent_state.params,
                         minibatch.obs,
                         minibatch.actions,
@@ -82,15 +50,82 @@ def update_ppo(
                         minibatch.advantages,
                         minibatch.returns,
                     )
+                )
                 agent_state = agent_state.apply_gradients(grads=grads)
-            return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
+                return agent_state, (
+                    loss,
+                    pg_loss,
+                    v_loss,
+                    entropy_loss,
+                    approx_kl,
+                    grads,
+                )
 
-        agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
+            agent_state, metrics = jax.lax.scan(
                 update_minibatch, agent_state, shuffled_storage
             )
-        return (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
+            return (agent_state, key), metrics
 
-    (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
+        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = (
+            jax.lax.scan(
                 update_epoch, (agent_state, key), (), length=args.update_epochs
             )
+        )
         return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
+
+
+"""
+Should be ok to use partial here, since the references to network,
+actor and critic should not change at runtime
+The cost of seperating concerns is to somehow pass these values
+that are now not in the same scope
+Chose to pass the actual apply functions, since they should never change
+Other option was to pass the networks, but how does it influence compilation when their
+class values would ever change? Better safe than sorry.
+"""
+
+
+@partial(jax.jit, static_argnums=(0, 1, 2))
+def get_action_and_value2(
+    network_apply,
+    actor_apply,
+    critic_apply,
+    params: flax.core.FrozenDict,
+    x: jnp.ndarray,
+    action: jnp.ndarray,
+):
+    hidden = network_apply(params["network_params"], x)
+    mean, log_std = actor_apply(params["actor_params"], hidden)
+    std = jnp.exp(log_std)
+
+    logprob = -0.5 * (
+        ((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)
+    ).sum(-1)
+    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
+    value = critic_apply(params["critic_params"], hidden).squeeze(-1)
+
+    return logprob, entropy, value
+
+
+def ppo_loss(
+    params, x, a, logp, mb_advantages, mb_returns, args, network, actor, critic
+):
+    newlogprob, entropy, newvalue = get_action_and_value2(
+        network.apply, actor.apply, critic.apply, params, x, a
+    )
+    logratio = newlogprob - logp
+    ratio = jnp.exp(logratio)
+    approx_kl = ((ratio - 1) - logratio).mean()
+
+    if args.norm_adv:
+        mb_advantages = (mb_advantages - mb_advantages.mean()) / (
+            mb_advantages.std() + 1e-8
+        )
+
+    pg_loss1 = -mb_advantages * ratio
+    pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
+    pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
+    v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
+    entropy_loss = entropy.mean()
+    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
+    return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))