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

fix(cleanup): ruff format

Browse source
This commit is contained in:
cedric 2026-04-09 22:13:07 +00:00
parent 71349af5b9
commit bf6434ea6c
1 changed files with 10 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -24,11 +24,13 @@ from brittle_star_project.MLPs.mlps import (
)
from brittle_star_project.ppo import PPO
def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float:
var_returns = jnp.var(returns)
explained_var = 1.0 - jnp.var(returns - values) / (var_returns + 1e-8)
return float(explained_var)
@jax.jit
def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
@ -412,25 +414,21 @@ class PPOTrainer:
jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item()
)
explained_var = _compute_explained_variance(
storage.values, storage.returns
)
explained_var = _compute_explained_variance(storage.values, storage.returns)
terminated = next_env_state.terminated # (num_envs,)
truncated = next_env_state.truncated # (num_envs,)
episode_lengths = self.episode_stats.returned_episode_lengths
num_terminated = int(jnp.sum(terminated).item())
num_truncated = int(jnp.sum(truncated).item())
avg_terminated_length = (
jnp.sum(episode_lengths * terminated) /
jnp.maximum(jnp.sum(terminated), 1)
num_truncated = int(jnp.sum(truncated).item())
avg_terminated_length = jnp.sum(episode_lengths * terminated) / jnp.maximum(
jnp.sum(terminated), 1
)
avg_truncated_length = (
jnp.sum(episode_lengths * truncated) /
jnp.maximum(jnp.sum(truncated), 1)
avg_truncated_length = jnp.sum(episode_lengths * truncated) / jnp.maximum(
jnp.sum(truncated), 1
)
return (
@ -448,7 +446,7 @@ class PPOTrainer:
num_terminated=num_terminated,
num_truncated=num_truncated,
avg_terminated_length=avg_terminated_length,
avg_truncated_length=avg_truncated_length
avg_truncated_length=avg_truncated_length,
),
)