From 6e942d896bfa38487065c6f4453c2ccec48803e2 Mon Sep 17 00:00:00 2001 From: JibrilExe Date: Fri, 10 Apr 2026 08:23:34 +0200 Subject: [PATCH] fix(gae): fixes #31 --- src/brittle_star_project/trainers/PPOTrainer.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 298e41c..a9e9b1b 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -175,11 +175,19 @@ def _compute_gae_once(carry, inp, gamma, gae_lambda): # jit applied on partial-wrapped wrapper method self._compute_gae_jit def _compute_gae_jit( - agent_state, storage, next_obs, next_done, gamma, gae_lambda, num_envs, sensor, critic + agent_state, + storage, + next_obs, + next_done, + gamma, + gae_lambda, + num_envs, + feature_extractor, + critic, ): next_value = critic.apply( agent_state.params["critic_params"], - sensor.apply(agent_state.params["sensor_params"], next_obs), + feature_extractor.apply(agent_state.params["feature_extractor_params"], next_obs), ).squeeze(-1) advantages = jnp.zeros((num_envs,)) @@ -244,7 +252,7 @@ class PPOTrainer: num_envs=self.args.num_envs, gamma=self.args.gamma, gae_lambda=self.args.gae_lambda, - sensor=self.sensor, + feature_extractor=self.feature_extractor, critic=self.critic, ) )