Merge pull request #9 from SELab-3-2026/chore/ppo

src/ppo.py

@@ -0,0 +1,161 @@
+from functools import partial
+
+import flax
+import jax
+import jax.numpy as jnp
+
+
+# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
+# with our need to seperate concerns
+class PPO:
+    def __init__(
+        self, args, input_network, action_network, critic, critic_network, message_passer=None
+    ):
+        self.args = args
+
+        if not message_passer:
+            message_passer = identity
+
+        self.ppo_loss_grad_fn = jax.value_and_grad(
+            partial(
+                ppo_loss,
+                args=args,
+                input_network_apply=input_network.apply,
+                action_network_apply=action_network.apply,
+                critic_apply=critic.apply,
+                critic_network_apply=critic_network.apply,
+                message_passer=message_passer,
+            ),
+            has_aux=True,
+        )
+
+    # This PPO class should be initialized only once,
+    # 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 update_epoch(carry, _):
+            agent_state, key = carry
+            key, subkey = jax.random.split(key)
+
+            def flatten(x):
+                return x.reshape((-1,) + x.shape[2:])
+
+            def convert_data(x):
+                x = jax.random.permutation(subkey, x)
+                return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
+
+            flatten_storage = jax.tree.map(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(
+                    agent_state.params,
+                    minibatch.obs,
+                    minibatch.actions,
+                    minibatch.logprobs,
+                    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,
+                )
+
+            agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
+            return (agent_state, key), metrics
+
+        (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
+"""
+
+
+@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
+def get_action_and_value2(
+    input_apply,
+    action_apply,
+    message_passer,
+    critic_apply,
+    critic_network_apply,
+    params: flax.core.FrozenDict,
+    x: jnp.ndarray,
+    action: jnp.ndarray,
+):
+    hidden_network = input_apply(params["network_params"], x)
+    hidden_critic = critic_network_apply(params["critic_network_params"], x)
+    hidden_network = message_passer(hidden_network)
+    mean, log_std = action_apply(params["actor_params"], hidden_network)
+    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_critic).squeeze(-1)
+
+    return logprob, entropy, value
+
+
+def ppo_loss(
+    params,
+    x,
+    a,
+    logp,
+    mb_advantages,
+    mb_returns,
+    args,
+    input_network_apply,
+    action_network_apply,
+    message_passer,
+    critic_apply,
+    critic_network_apply,
+):
+    newlogprob, entropy, newvalue = get_action_and_value2(
+        input_network_apply,
+        action_network_apply,
+        message_passer,
+        critic_apply,
+        critic_network_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))
+
+
+def identity(hidden):
+    """
+    Used for seamless jax integration,
+    avoids having branching inside jitted function,
+    used as message_passer in case it is not given,
+    (in case of centralized lvl)
+    """
+
+    return hidden