diff --git a/configs/hpc/wandb_expand.yaml b/configs/hpc/wandb_expand.yaml new file mode 100644 index 0000000..432160f --- /dev/null +++ b/configs/hpc/wandb_expand.yaml @@ -0,0 +1,10 @@ +exp_name: "explained_var_fun_more_steps" +seed: 42 +track: true +wandb_project_name: "LET-THERE-BE-MORE-LOGGING" +wandb_entity: "SEL3-2026-Groep-4" + +num_envs: 16 +num_steps: 256 +total_timesteps: 50000 +cuda: true \ No newline at end of file diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 432030a..1e554f8 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -28,6 +28,12 @@ from brittle_star_project.ppo import PPO def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: return jnp.clip(action, low, high) +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 @@ -41,7 +47,6 @@ def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: ) -# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit def _get_action_and_value_noise( sensor: GenericDenseLayersWithActivation, feature_extractor: GenericDenseLayersWithActivation, @@ -58,7 +63,6 @@ def _get_action_and_value_noise( agent_state.params["feature_extractor_params"], next_obs ) - # Continuous actions: sample from a Gaussian parameterized by the actor mean, log_std = actor.apply(agent_state.params["actor_params"], hidden) key, subkey = jax.random.split(key) noise = jax.random.normal(subkey, shape=mean.shape) @@ -74,7 +78,6 @@ def _get_action_and_value_noise( return clipped_action, logprob, value.squeeze(-1), key -# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit def _step_once( carry, _, @@ -108,15 +111,13 @@ def _step_once( return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage -# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): next_env_state = env_step_fn(env_state, action) - # Extract per-environment signals from the state object - reward = next_env_state.reward # (num_envs,) - terminated = next_env_state.terminated # (num_envs,) - truncated = next_env_state.truncated # (num_envs,) - done = terminated | truncated # (num_envs,) + reward = next_env_state.reward + terminated = next_env_state.terminated + truncated = next_env_state.truncated + done = terminated | truncated new_episode_return = episode_stats.episode_returns + reward new_episode_length = episode_stats.episode_lengths + 1 @@ -138,7 +139,6 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn): ) -# jit applied in wrapper method self._rollout_jit using partial def _rollout_jit( agent_state, episode_stats, @@ -173,7 +173,6 @@ def _rollout_jit( return agent_state, episode_stats, next_obs, next_done, storage, key, env_state -# removed jit: used in _compute_gae_jit, so will be compiled with _compute_gae_jit def _compute_gae_once(carry, inp, gamma, gae_lambda): advantages = carry nextdone, nextvalues, curvalues, reward = inp @@ -183,13 +182,20 @@ def _compute_gae_once(carry, inp, gamma, gae_lambda): return advantages, advantages -# 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,)) @@ -205,14 +211,18 @@ def _compute_gae_jit( @dataclass -class LossInfo: - # todo: better typing - loss: Any - pg_loss: Any - v_loss: Any - entropy_loss: Any - approx_kl: Any - avg_episodic_return: Any +class TrainingMeasurements: + loss: jnp.ndarray + pg_loss: jnp.ndarray + v_loss: jnp.ndarray + entropy_loss: jnp.ndarray + approx_kl: jnp.ndarray + avg_episodic_return: float + explained_variance: float + num_terminated: int + num_truncated: int + avg_terminated_length: Any + avg_truncated_length: Any class PPOTrainer: @@ -253,7 +263,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, ) ) @@ -277,11 +287,8 @@ class PPOTrainer: sensor = GenericDenseLayersWithActivation() feature_extractor = GenericDenseLayersWithActivation() - actor = Actor( - action_dim=self.env.single_action_space.shape[0] - ) # continuous actions for MJX + actor = Actor(action_dim=self.env.single_action_space.shape[0]) critic = OneDenseLayerMLP() - # messenger = OneDenseLayerMLP() return sensor, feature_extractor, actor, critic def _init_agent_state(self) -> TrainState: @@ -332,7 +339,7 @@ class PPOTrainer: def _init_episode_stats(self) -> EpisodeStatistics: self.logger.info("[EPISODE STATS]: Initializing episode stats...") - return EpisodeStatistics( # type: ignore[call-arg] + return EpisodeStatistics( episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32), episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32), returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32), @@ -363,21 +370,26 @@ class PPOTrainer: episode_stats, start_time, iteration_time_start, - loss_info, + training_measurements, ): metrics = { - "charts/avg_episodic_return": loss_info.avg_episodic_return, + "charts/avg_episodic_return": training_measurements.avg_episodic_return, "charts/avg_episodic_length": np.mean( jax.device_get(episode_stats.returned_episode_lengths) ), "charts/learning_rate": self.agent_state.opt_state[1] .hyperparams["learning_rate"] .item(), - "losses/value_loss": loss_info.v_loss[-1, -1].item(), - "losses/policy_loss": loss_info.pg_loss[-1, -1].item(), - "losses/entropy": loss_info.entropy_loss[-1, -1].item(), - "losses/approx_kl": loss_info.approx_kl[-1, -1].item(), - "losses/loss": loss_info.loss[-1, -1].item(), + "charts/explained_variance": training_measurements.explained_variance, + "charts/num_terminated": training_measurements.num_terminated, + "charts/num_truncated": training_measurements.num_truncated, + "charts/avg_terminated_ep_length": training_measurements.avg_terminated_length, + "charts/avg_truncated_ep_length": training_measurements.avg_truncated_length, + "losses/value_loss": training_measurements.v_loss[-1, -1].item(), + "losses/policy_loss": training_measurements.pg_loss[-1, -1].item(), + "losses/entropy": training_measurements.entropy_loss[-1, -1].item(), + "losses/approx_kl": training_measurements.approx_kl[-1, -1].item(), + "losses/loss": training_measurements.loss[-1, -1].item(), "charts/SPS": int(global_step / (time.time() - start_time)), "charts/SPS_update": int( self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start) @@ -418,17 +430,39 @@ class PPOTrainer: jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item() ) + explained_var = _compute_explained_variance(storage.values, storage.returns) + + terminated = next_env_state.terminated + truncated = next_env_state.truncated + 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 + ) + + avg_truncated_length = jnp.sum(episode_lengths * truncated) / jnp.maximum( + jnp.sum(truncated), 1 + ) + return ( next_env_state, next_obs, next_done, - LossInfo( + TrainingMeasurements( loss=loss, pg_loss=pg_loss, v_loss=v_loss, entropy_loss=entropy_loss, approx_kl=approx_kl, avg_episodic_return=avg_episodic_return, + explained_variance=explained_var, + num_terminated=num_terminated, + num_truncated=num_truncated, + avg_terminated_length=avg_terminated_length, + avg_truncated_length=avg_truncated_length, ), ) @@ -473,12 +507,18 @@ class PPOTrainer: for iteration in iter_bar: iteration_time_start = time.time() - env_state, next_obs, next_done, loss_info = self._step( + env_state, next_obs, next_done, training_measurements = self._step( env_state, next_obs, next_done, iteration=iteration ) global_step += self.args.num_steps * self.args.num_envs - self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info) + self._log( + global_step, + self.episode_stats, + start_time, + iteration_time_start, + training_measurements, + ) sps = int(global_step / (time.time() - start_time)) remaining_steps = self.args.total_timesteps - global_step @@ -489,7 +529,7 @@ class PPOTrainer: f"Iteration {iteration}/{self.args.num_iterations} | " f"Step {global_step}/{self.args.total_timesteps} | " f"SPS {sps} | " - f"Return {loss_info.avg_episodic_return:.4f} | " + f"Return {training_measurements.avg_episodic_return:.4f} | " f"ETA {eta_str}" )