1
Fork 0

fix: formatting + linting

This commit is contained in:
Robin Meersman 2026-04-06 11:55:05 +02:00
parent 3c7670e1b4
commit 9c5203fc7d
3 changed files with 6 additions and 4 deletions

View file

@ -9,7 +9,6 @@ import jax
import jax.numpy as jnp
import numpy as np
import optax
import torch
import tqdm
from flax.training.train_state import TrainState
from torch.utils.tensorboard import SummaryWriter
@ -381,7 +380,9 @@ class PPOTrainer:
self._ppo.update_ppo(self.agent_state, storage, self.key)
)
avg_episodic_return = float(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 (
next_env_state,