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

fix(PPOTrainer.py): finished all jit-related bugs

Browse source
This commit is contained in:
Robin Meersman 2026-04-06 11:53:22 +02:00
parent 62e11ce243
commit ca17574272
1 changed files with 30 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

62
experiments/PPOTrainer.py
View file

@ -39,7 +39,7 @@ def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
) )
@jax.jit # removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
def _get_action_and_value_noise( def _get_action_and_value_noise(
sensor: GenericDenseLayersWithActivation, sensor: GenericDenseLayersWithActivation,
feature_extractor: GenericDenseLayersWithActivation, feature_extractor: GenericDenseLayersWithActivation,
@ -65,7 +65,7 @@ def _get_action_and_value_noise(
return action, logprob, value.squeeze(-1), key return action, logprob, value.squeeze(-1), key
@jax.jit # removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
def _step_once( def _step_once(
carry, carry,
_, _,
@ -97,8 +97,8 @@ def _step_once(
return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
@jax.jit # removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
def _step_env_wrapped(env_step_fn, env_state, action, episode_stats): def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
next_env_state = env_step_fn(env_state, action) next_env_state = env_step_fn(env_state, action)
# Extract per-environment signals from the state object # Extract per-environment signals from the state object
@ -127,10 +127,7 @@ def _step_env_wrapped(env_step_fn, env_state, action, episode_stats):
) )
@partial( # jit applied in wrapper method self._rollout_jit using partial
jax.jit,
static_argnames=("env_step_fn", "sensor", "feature_extractor", "actor", "critic"),
)
def _rollout_jit( def _rollout_jit(
agent_state, agent_state,
episode_stats, episode_stats,
@ -139,7 +136,7 @@ def _rollout_jit(
next_done, next_done,
key, key,
max_steps, max_steps,
env_step_fn, step_env_fn,
sensor: GenericDenseLayersWithActivation, sensor: GenericDenseLayersWithActivation,
feature_extractor: GenericDenseLayersWithActivation, feature_extractor: GenericDenseLayersWithActivation,
actor: Actor, actor: Actor,
@ -152,7 +149,7 @@ def _rollout_jit(
feature_extractor=feature_extractor, feature_extractor=feature_extractor,
actor=actor, actor=actor,
critic=critic, critic=critic,
env_step_fn=partial(_step_env_wrapped, env_step_fn=env_step_fn), env_step_fn=step_env_fn,
), ),
(agent_state, episode_stats, next_obs, next_done, key, env_state), (agent_state, episode_stats, next_obs, next_done, key, env_state),
(), (),
@ -161,7 +158,7 @@ def _rollout_jit(
return agent_state, episode_stats, next_obs, next_done, storage, key, env_state return agent_state, episode_stats, next_obs, next_done, storage, key, env_state
@jax.jit # removed jit: used in _compute_gae_jit, so will be compiled with _compute_gae_jit
def _compute_gae_once(carry, inp, gamma, gae_lambda): def _compute_gae_once(carry, inp, gamma, gae_lambda):
advantages = carry advantages = carry
nextdone, nextvalues, curvalues, reward = inp nextdone, nextvalues, curvalues, reward = inp
@ -171,12 +168,9 @@ def _compute_gae_once(carry, inp, gamma, gae_lambda):
return advantages, advantages return advantages, advantages
@partial( # jit applied on partial-wrapped wrapper method self._compute_gae_jit
jax.jit,
static_argnames=("sensor", "critic"),
)
def _compute_gae_jit( def _compute_gae_jit(
agent_state, storage, next_obs, next_done, sensor, critic, gamma, gae_lambda, num_envs agent_state, storage, next_obs, next_done, gamma, gae_lambda, num_envs, sensor, critic
): ):
next_value = critic.apply( next_value = critic.apply(
agent_state.params["critic_params"], agent_state.params["critic_params"],
@ -221,17 +215,26 @@ class PPOTrainer:
self.actor.apply = jax.jit(self.actor.apply) self.actor.apply = jax.jit(self.actor.apply)
self.critic.apply = jax.jit(self.critic.apply) self.critic.apply = jax.jit(self.critic.apply)
self._rollout_jit = partial( self._rollout_jit = jax.jit(
_rollout_jit, partial(
sensor=self.sensor, _rollout_jit,
feature_extractor=self.feature_extractor, max_steps=self.args.num_steps,
actor=self.actor, step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
critic=self.critic, sensor=self.sensor,
feature_extractor=self.feature_extractor,
actor=self.actor,
critic=self.critic,
)
) )
self._compute_gae_jit = partial( self._compute_gae_jit = jax.jit(
_compute_gae_jit, partial(
sensor=self.sensor, _compute_gae_jit,
critic=self.critic, num_envs=self.args.num_envs,
gamma=self.args.gamma,
gae_lambda=self.args.gae_lambda,
sensor=self.sensor,
critic=self.critic,
)
) )
self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor) self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)
@ -315,8 +318,6 @@ class PPOTrainer:
next_obs, next_obs,
next_done, next_done,
self.key, self.key,
self.args.num_steps,
self.env.step,
) )
def _compute_gae(self, storage, next_obs, next_done) -> Storage: def _compute_gae(self, storage, next_obs, next_done) -> Storage:
@ -325,9 +326,6 @@ class PPOTrainer:
storage, storage,
next_obs, next_obs,
next_done, next_done,
self.args.gamma,
self.args.gae_lambda,
self.args.num_envs,
) )
def _log( def _log(
@ -383,7 +381,7 @@ class PPOTrainer:
self._ppo.update_ppo(self.agent_state, storage, self.key) self._ppo.update_ppo(self.agent_state, storage, self.key)
) )
avg_episodic_return = jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)) avg_episodic_return = float(jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)))
return ( return (
next_env_state, next_env_state,