tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

fix(gae): fixes #31

Browse source
This commit is contained in:
JibrilExe 2026-04-10 08:23:34 +02:00
parent bf6434ea6c
commit 6e942d896b
1 changed files with 11 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

14
src/brittle_star_project/trainers/PPOTrainer.py
View file

@ -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,
)
)