From 4e3cd0c6738b0efd128c70d170287965cb58aa71 Mon Sep 17 00:00:00 2001
From: Robin Meersman <echteenrobin@gmail.com>
Date: Thu, 7 May 2026 14:07:19 +0200
Subject: [PATCH] fix(ppo): identity function now takes 2 params, like
 message_passer

---
 src/brittle_star_project/ppo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py
index fffc26b..fb27ccd 100644
--- a/src/brittle_star_project/ppo.py
+++ b/src/brittle_star_project/ppo.py
@@ -190,7 +190,7 @@ def ppo_loss(
     return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
 
 
-def identity(hidden):
+def identity(_, hidden):
     """
     Used for seamless jax integration,
     avoids having branching inside jitted function,