import datetime import random import time from dataclasses import asdict, dataclass from functools import partial from typing import Any, Optional import jax import jax.numpy as jnp import numpy as np import optax import flax.linen as nn from flax.training.train_state import TrainState from experiment_logger import get_logger from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.dataclasses import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from brittle_star_project.MLPs.mlps import ( Actor, AgentParams, GenericDenseLayersWithActivation, MessagePasser, OneDenseLayerMLP, Storage, ) from brittle_star_project.ppo import PPO from brittle_star_project.environment import MorphMode from brittle_star_project.utils import logged_jit logger11 = get_logger() # TODO: move to config _ALLOWED_OBS_KEYS = { "joint_position", "joint_velocity", "joint_actuator_force", "actuator_force", "disk_position", "disk_rotation", "disk_linear_velocity", "disk_angular_velocity", "unit_xy_direction_to_target", "xy_distance_to_target", } # TODO: clip scaled reward? def build_adjacency(segments_per_arm, mode: MorphMode): num_arms = sum(1 for s in segments_per_arm if s > 0) num_segments = sum(segments_per_arm) # FOR NOW SEMI HARDCODE: # CENTRALIZED: 1 agent, no stress, adja = 1,1 = [[1]] # FULLY CONNECTED: 5 agents: adj = alle 1 # CENTRAL DISK:#arms= 5 agents, only neighbor as adjacent so diagonal kinda.. # ARM = #segments agents: diago kinda, but extra, center ring too, put center mlps first or.. if mode == MorphMode.CENTRALIZED: return jnp.ones((1, 1)) if mode == MorphMode.FULLY_CONNECTED: adj = jnp.ones((num_arms, num_arms)) # everybody adjacent everybody return adj if mode == MorphMode.RING: # ring adj = jnp.zeros((num_arms, num_arms)) for i in range(num_arms): adj = adj.at[i, i].set(1) # self adj = adj.at[i, (i - 1) % num_arms].set(1) adj = adj.at[i, (i + 1) % num_arms].set(1) # left and right.. return adj if mode == MorphMode.SEGMENT: num_nodes = num_arms + num_segments adj = jnp.zeros((num_nodes, num_nodes)) # first ring for i in range(num_arms): # self adj = adj.at[i, i].set(1) # ring neighbors adj = adj.at[i, (i - 1) % num_arms].set(1) adj = adj.at[i, (i + 1) % num_arms].set(1) # then segment chains idx = 0 for arm_idx, seg_count in enumerate(segments_per_arm): for i in range(seg_count): seg_node = num_arms + idx + i adj = adj.at[seg_node, seg_node].set(1) if i > 0: adj = adj.at[seg_node, seg_node - 1].set(1) if i < seg_count - 1: adj = adj.at[seg_node, seg_node + 1].set(1) idx += seg_count idx = 0 for arm_idx, seg_count in enumerate(segments_per_arm): first_seg = num_arms + idx # first segment of this arm # connect ring node first segment adj = adj.at[arm_idx, first_seg].set(1) adj = adj.at[first_seg, arm_idx].set(1) idx += seg_count return adj @logged_jit 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) @logged_jit def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate): frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations return learning_rate * frac @logged_jit def _normalize_obs(obs, mean, var, eps=1e-8): return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0) def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, segments_per_arm: jnp.ndarray): num_segments = int(segments_per_arm.sum()) num_arms = int(jnp.where(segments_per_arm > 0, 1, 0).sum()) @logged_jit def _filter_and_flatten(o) -> jnp.ndarray: # vmap feeds one env at a time — v has NO batch dim here # shapes are e.g. (n_features,) or (n_nodes, feat) values = [] for key in sorted(o.keys()): if key not in _ALLOWED_OBS_KEYS: continue v = o[key] if v.size == 0: continue # -------- CENTRALIZED -------- if morph_mode == MorphMode.CENTRALIZED: values.append(v.reshape(1, -1)) # (1, feat) continue # -------- SPLIT TO SEGMENTS -------- if key in _JOINT_SCALED_KEYS: if morph_mode == MorphMode.SEGMENT: center_size = num_arms * 3 * 2 v_center = v[:center_size].reshape(num_arms, 3 * 2) # (arms, 6) v_segs = v[center_size:].reshape(-1, 2) # (segs, 2) values.append(jnp.concatenate([v_center, v_segs], axis=0)) # (arms+segs, ?) continue v = v.reshape(num_arms, -1) # (n_arms, 2) elif key in _SEGMENT_SCALED_KEYS: v = v[:, None] # (segments, 1) else: # global key, broadcast to all nodes n_nodes = (num_segments + num_arms) if morph_mode == MorphMode.SEGMENT else num_arms v = jnp.repeat(v[None, :], n_nodes, axis=0) # (n_nodes, feat) # -------- SEGMENT MODE -------- if morph_mode == MorphMode.SEGMENT: values.append(v) # (n_nodes, feat) continue # -------- ARM MODE -------- v = v.reshape(num_arms, -1) values.append(v) # (n_arms, feat) return jnp.concatenate(values, axis=-1) # (n_nodes, total_feat) return jax.vmap(_filter_and_flatten)(obs_dict) # output: (batch, n_nodes, total_feat) # Observation keys whose size scales with the number of joints (2 per segment). _JOINT_SCALED_KEYS = frozenset( { # TODO CODE SMELL "joint_position", "joint_velocity", "joint_actuator_force", "actuator_force", } ) # Observation keys whose size scales with the number of segments (1 per segment). _SEGMENT_SCALED_KEYS = frozenset( { "segment_contact", } ) # TODO: update to work with extra dimension + message passing def _get_action_and_value_noise( sensor: nn.Module, feature_extractor: nn.Module, actor: nn.Module, critic: nn.Module, message_passer: Optional[nn.Module], agent_state: TrainState, next_obs: jnp.ndarray, key, action_low, action_high, adj_matrix: jnp.ndarray, ): # (B, n_nodes, feat) hidden = apply_per_node(sensor, agent_state.params["sensor_params"], next_obs) get_logger().debug(f"[_get_action_and_value_noise] hidden (before): {hidden.shape}") if message_passer is not None: params = agent_state.params["message_passer_params"] # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension hidden = jax.vmap(lambda x: message_passer.apply(params, x, adj_matrix))(hidden) get_logger().debug(f"[_get_action_and_value_noise] hidden (after): {hidden.shape}") hidden_critic = apply_shared( feature_extractor, agent_state.params["feature_extractor_params"], next_obs ) mean, log_std = apply_per_node(actor, agent_state.params["actor_params"], hidden) log_std = jnp.clip(log_std, -5, 2) key, subkey = jax.random.split(key) noise = jax.random.normal(subkey, shape=mean.shape) std = jnp.exp(log_std) raw_action = mean + noise * std clipped_action = _clip_action(raw_action, action_low, action_high) flat_clipped_action = clipped_action.reshape( clipped_action.shape[0], -1 ) # concat the per agent, keep the envs dim (batch, agent * action) logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum( axis=(-2, -1) ) value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic) return flat_clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key # TODO: update to work vectorized (sensor, actor, message passer) + message passing def _step_once( carry, _, env_step_fn, adj_matrix, sensor: nn.Module, feature_extractor: nn.Module, actor: nn.Module, critic: nn.Module, message_passer: Optional[nn.Module], action_low, action_high, ): agent_state, episode_stats, obs, done, key, env_state = carry flat_clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( sensor, feature_extractor, actor, critic, message_passer, agent_state, obs, key, action_low, action_high, adj_matrix, ) logger11.debug(f"[_step_once] raw_action: {raw_action.shape}") logger11.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}") # Supporting signals (often where mismatch originates) logger11.debug(f"[_step_once] logprob: {logprob.shape}") logger11.debug(f"[_step_once] value: {value.shape}") logger11.debug(f"[_step_once] mean: {mean.shape}") logger11.debug(f"[_step_once] std: {std.shape}") # ---- ENV STEP ---- episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn( episode_stats, env_state, flat_clipped_action ) logger11.debug(f"[_step_once] next_obs: {next_obs.shape}") logger11.debug(f"[_step_once] reward: {reward.shape}") logger11.debug(f"[_step_once] next_done: {next_done.shape}") storage = Storage( obs=obs, actions=raw_action, raw_actions=raw_action, logprobs=logprob, dones=done, values=value, rewards=reward, means=mean, stds=std, returns=jnp.zeros_like(reward), advantages=jnp.zeros_like(reward), ) return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage def _reward_fn(env_state, next_env_state): # if delta distance positive ==> brittle star walking away from target delta_distance = ( next_env_state.observations["xy_distance_to_target"] - env_state.observations["xy_distance_to_target"] ).squeeze(-1) env_reward = next_env_state.reward clipped_env_reward = jnp.clip(100 * env_reward, -10, 10) time_penalty = 0.1 distance_penalty = jnp.clip(0.5 * delta_distance, -0.5, 0.5) penalty = time_penalty + distance_penalty return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty) def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, morph_mode, segments_per_arm): next_env_state = env_step_fn(env_state, action) reward = _reward_fn(env_state, next_env_state) 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 episode_stats = episode_stats.replace( episode_returns=new_episode_return * (1 - done), episode_lengths=new_episode_length * (1 - done), returned_episode_returns=jnp.where( done, new_episode_return, episode_stats.returned_episode_returns ), returned_episode_lengths=jnp.where( done, new_episode_length, episode_stats.returned_episode_lengths ), ) return ( episode_stats, next_env_state, ( _convert_obs_dict_to_array_morphology( next_env_state.observations, morph_mode, segments_per_arm ), reward, done, ), ) def apply_per_node(net, params, x): # params: (nodes, ...) # x: (batch, nodes, feat) def apply_single_node(p, x_node): # x_node: (batch, feat) return jax.vmap(lambda xi: net.apply(p, xi))(x_node) return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x) def apply_shared(net, params, x): # x: (batch, nodes, feat) # If the critic expects a single vector per environment: batch_size = x.shape[0] x_flattened = x.reshape(batch_size, -1) return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened) # TODO: update to work with extra dimension + message passing def _rollout_jit( agent_state, episode_stats, env_state, next_obs, next_done, key, max_steps, step_env_fn, sensor: nn.Module, feature_extractor: nn.Module, actor: nn.Module, critic: nn.Module, message_passer: Optional[nn.Module], action_low, action_high, adj_matrix, ): (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan( partial( _step_once, sensor=sensor, feature_extractor=feature_extractor, actor=actor, critic=critic, message_passer=message_passer, env_step_fn=step_env_fn, action_low=action_low, action_high=action_high, adj_matrix=adj_matrix, ), (agent_state, episode_stats, next_obs, next_done, key, env_state), (), max_steps, ) return agent_state, episode_stats, next_obs, next_done, storage, key, env_state def _compute_gae_once(carry, inp, gamma, gae_lambda): advantages = carry nextdone, nextvalues, curvalues, reward = inp nextnonterminal = 1.0 - nextdone delta = reward + gamma * nextvalues * nextnonterminal - curvalues advantages = delta + gamma * gae_lambda * nextnonterminal * advantages return advantages, advantages def _compute_gae_jit( agent_state, storage, next_obs, next_done, gamma, gae_lambda, num_envs, feature_extractor, critic, ): next_value = apply_shared( critic, agent_state.params["critic_params"], apply_shared(feature_extractor, agent_state.params["feature_extractor_params"], next_obs), ).squeeze(-1) advantages = jnp.zeros((num_envs,)) dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0) values = jnp.concatenate([storage.values, next_value[None, :]], axis=0) _, advantages = jax.lax.scan( partial(_compute_gae_once, gamma=gamma, gae_lambda=gae_lambda), advantages, (dones[1:], values[1:], values[:-1], storage.rewards), reverse=True, ) returns = advantages + storage.values advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) return storage.replace(advantages=advantages, returns=returns) @dataclass 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: def __init__( self, cfg: BrittleStarConfig, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str, ): self.cfg = cfg self.ppo = cfg.ppo self.experiment = cfg.experiment self.logging_cfg = cfg.logging self.env = env self.run_dir = run_dir self.run_name = run_name self.logger = get_logger() # Derived runtime fields self.batch_size = self.ppo.num_envs * self.ppo.num_steps self.num_iterations = self.ppo.total_timesteps // self.batch_size self.key = jax.random.PRNGKey(self.experiment.seed) self.morph_mode = self.cfg.morphology.morph_mode self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32) self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}") self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode) ( self.sensor, self.message_passer, self.actor, self.feature_extractor, self.critic, self.needed_copies, ) = self._init_agent() self.sensor.apply = logged_jit(self.sensor.apply) self.feature_extractor.apply = logged_jit(self.feature_extractor.apply) self.actor.apply = logged_jit(self.actor.apply) self.critic.apply = logged_jit(self.critic.apply) action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) self._rollout_jit = logged_jit( partial( _rollout_jit, max_steps=self.ppo.num_steps, step_env_fn=partial( _step_env_wrapped, env_step_fn=self.env.step, morph_mode=self.morph_mode, segments_per_arm=self.segments_per_arm, ), sensor=self.sensor, feature_extractor=self.feature_extractor, actor=self.actor, critic=self.critic, message_passer=self.message_passer, action_low=action_low, action_high=action_high, adj_matrix=self.adj, ) ) self._compute_gae_jit = logged_jit( partial( _compute_gae_jit, num_envs=self.ppo.num_envs, gamma=self.ppo.gamma, gae_lambda=self.ppo.gae_lambda, feature_extractor=self.feature_extractor, critic=self.critic, ) ) def apply_sensor(p, x): return apply_per_node(self.sensor, p, x) def apply_actor(p, x): return apply_per_node(self.actor, p, x) def apply_critic(p, x): return apply_shared(self.critic, p, x) def apply_feature(p, x): return apply_shared(self.feature_extractor, p, x) def apply_message_passer(p, x): return apply_shared(self.message_passer, p, x) self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature) self.agent_state = self._init_agent_state() self.episode_stats = self._init_episode_stats() self._init_random() def _init_random(self): self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}") random.seed(self.experiment.seed) np.random.seed(self.experiment.seed) def _init_agent(self): self.logger.info("[AGENT]: Initializing agent...") match self.morph_mode: case MorphMode.CENTRALIZED: needed_copies = 1 case MorphMode.FULLY_CONNECTED | MorphMode.RING: needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() case MorphMode.SEGMENT: needed_copies = ( self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum() ).item() actor = Actor(action_dim=self.env.single_action_space.shape[0]) sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) message_passer: Optional[nn.Module] = ( MessagePasser( hidden_dim=300, num_propagation_steps=self.cfg.architecture.message_passing_steps or 4, ) if self.morph_mode != MorphMode.CENTRALIZED else None ) feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) critic = OneDenseLayerMLP() return sensor, message_passer, actor, feature_extractor, critic, needed_copies def _init_agent_state(self) -> TrainState: self.logger.info("[AGENT STATE]: Initializing agent state...") self.key, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = ( jax.random.split(self.key, 6) ) dummy_reset = self.env.reset(seed=0) for k, v in dummy_reset.observations.items(): self.logger.debug(k, v.shape) sample_obs = _convert_obs_dict_to_array_morphology( dummy_reset.observations, self.morph_mode, self.segments_per_arm, )[0] # take first env self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}") self.obs_mean = jnp.zeros((len(sample_obs),)) self.obs_var = jnp.ones((len(sample_obs),)) self.obs_count = 1e-4 self.logger.debug(f"[_init_agent_state] obs_mean: {self.obs_mean.shape}") self.logger.debug(f"[_init_agent_state] obs_var: {self.obs_var.shape}") self.logger.debug(f"[_init_agent_state]: Needed copies: {self.needed_copies}") sensor_keys = jax.random.split(sensor_key, self.needed_copies) actor_keys = jax.random.split(actor_key, self.needed_copies) # note: assumed only 1 message passer needed for now # message_passer_keys = jax.random.split(message_passer_key, self.needed_copies) # (needed_copies, 175) sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys) # self.logger.debug( # f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}" # ) single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params) # self.logger.debug( # f"[_init_agent_state] single_sensor_param: { # jax.tree.map(lambda x: x.shape, single_sensor_param) # }" # ) sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs) self.logger.debug( f"[_init_agent_state] sensor_params_sample shape: {sensor_params_sample.shape}" ) actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys) # self.logger.debug( # f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}" # ) message_passer_params = {} if self.morph_mode != MorphMode.CENTRALIZED: assert self.message_passer is not None, "MessagePasser is None" # message_passer_params = jax.vmap( # lambda k: self.message_passer.init( # k, self.sensor.apply(single_sensor_param, sample_obs), self.adj # ) # )(message_passer_keys) message_passer_params = self.message_passer.init( message_passer_key, self.sensor.apply(single_sensor_param, sample_obs), self.adj, ) # self.logger.debug( # f"[_init_agent_state] message_passer_params: { # jax.tree.map(lambda x: x.shape, message_passer_params) # }" # ) flat_obs = sample_obs.reshape(-1) # BECAUSE 1 centralized critic self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}") feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs) # self.logger.debug( # f"[_init_agent_state] feature_extractor_params: { # jax.tree.map(lambda x: x.shape, feature_extractor_params) # }" # ) critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs) self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}") critic_params = self.critic.init(critic_key, critic_input) # self.logger.debug( # f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}" # ) return TrainState.create( apply_fn=None, params=asdict( AgentParams( sensor_params, actor_params, critic_params, feature_extractor_params, message_passer_params, ) ), tx=optax.chain( optax.clip_by_global_norm(self.ppo.max_grad_norm), optax.inject_hyperparams(optax.adam)( learning_rate=partial( _linear_schedule, minibatch_count=self.ppo.num_minibatches, update_epochs=self.ppo.update_epochs, num_iterations=self.num_iterations, learning_rate=self.ppo.learning_rate, ) if self.ppo.anneal_lr else self.ppo.learning_rate, eps=1e-5, ), ), ) def _init_episode_stats(self) -> EpisodeStatistics: self.logger.info("[EPISODE STATS]: Initializing episode stats...") return EpisodeStatistics( episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32), episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32), returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), ) def _update_obs_stats(self, obs: jnp.ndarray): batch_mean = jnp.mean(obs, axis=0) batch_var = jnp.var(obs, axis=0) batch_count = obs.shape[0] delta = batch_mean - self.obs_mean total_count = self.obs_count + batch_count new_mean = self.obs_mean + delta * batch_count / total_count m_a = self.obs_var * self.obs_count m_b = batch_var * batch_count M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count new_var = M2 / total_count self.obs_mean = new_mean self.obs_var = new_var self.obs_count = total_count def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: return self._rollout_jit( self.agent_state, self.episode_stats, env_state, next_obs, next_done, self.key, ) def _compute_gae(self, storage, next_obs, next_done) -> Storage: return self._compute_gae_jit( self.agent_state, storage, next_obs, next_done, ) def _log( self, global_step, episode_stats, start_time, iteration_time_start, training_measurements, storage, ): data = jax.device_get( { "rewards": storage.rewards, "values": storage.values, "returns": storage.returns, "advantages": storage.advantages, } ) rollout_metrics = { "rollout/reward_mean": float(np.mean(data["rewards"])), "rollout/return_mean": float(np.mean(data["returns"])), "rollout/value_mean": float(np.mean(data["values"])), "rollout/advantage_mean": float(np.mean(data["advantages"])), "rollout/advantage_std": float(np.std(data["advantages"])), "rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)), } metrics = { "charts/episodic_return": training_measurements.avg_episodic_return, "charts/episodic_length": float( np.mean(jax.device_get(episode_stats.returned_episode_lengths)) ), "charts/explained_variance": training_measurements.explained_variance, "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(), "charts/learning_rate": self.agent_state.opt_state[1] .hyperparams["learning_rate"] .item(), "charts/SPS": int(global_step / (time.time() - start_time)), "charts/SPS_update": int( self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start) ), "termi_trunci/num_terminated": training_measurements.num_terminated, "termi_trunci/num_truncated": training_measurements.num_truncated, "termi_trunci/avg_terminated_ep_length": training_measurements.avg_terminated_length, "termi_trunci/avg_truncated_ep_length": training_measurements.avg_truncated_length, **rollout_metrics, } self.logger.log(metrics, step=global_step) def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple: if iteration == 1: self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}") self.logger.info(f"[_step] next_obs (in): {next_obs.shape}") ( self.agent_state, self.episode_stats, next_obs, next_done, storage, self.key, next_env_state, ) = self._rollout(env_state, next_obs, next_done) self.logger.info(f"[_step] next_obs (post-rollout): {next_obs.shape}") if iteration == 1: self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}") storage = self._compute_gae(storage, next_obs, next_done) self.logger.info(f"[_step] storage.obs (post-gae): {storage.obs.shape}") if iteration == 1: self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}") self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = ( self._ppo.update_ppo(self.agent_state, storage, self.key) ) if iteration == 1: self.logger.log_non_interactive(f"First PPO update completed: {time.ctime()}") avg_episodic_return = float( 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, 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, ), storage, ) def _close(self): self.env.close() def _save_model(self, model_path: str): self.logger.info("[SAVE]: Saving the final model...") self.logger.save_final_model(params=self.agent_state.params, metadata=asdict(self.cfg)) def _save_checkpoint(self, iteration: int): self.logger.info(f"[SAVE]: Saving checkpoint at iteration {iteration}...") self.logger.save_checkpoint( params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg) ) def train(self): """ Train the PPO agent for a specified number of iterations. Closes the environment at the end of training. """ self.logger.info(f"running name: {self.run_name}") self.logger.info("[TRAIN]: Resetting environment...") self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}") env_state = self.env.reset(seed=self.experiment.seed) next_obs = _convert_obs_dict_to_array_morphology( env_state.observations, self.morph_mode, self.segments_per_arm, ) self.logger.info(f"[train] next_obs: {next_obs.shape}") next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_) self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") global_step = 0 start_time = time.time() iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1)) for iteration in iter_bar: iteration_time_start = time.time() env_state, next_obs, next_done, training_measurements, storage = self._step( env_state, next_obs, next_done, iteration=iteration ) self.logger.info(f"[train] next_obs (post-step): {next_obs.shape}") self._update_obs_stats(next_obs) next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var) global_step += self.ppo.num_steps * self.ppo.num_envs self._log( global_step, self.episode_stats, start_time, iteration_time_start, training_measurements, storage, ) sps = int(global_step / (time.time() - start_time)) remaining_steps = self.ppo.total_timesteps - global_step eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 eta_str = str(datetime.timedelta(seconds=eta_seconds)) self.logger.log_non_interactive( f"Iteration {iteration}/{self.num_iterations} | " f"Step {global_step}/{self.ppo.total_timesteps} | " f"SPS {sps} | " f"Return {training_measurements.avg_episodic_return:.4f} | " f"ETA {eta_str}" ) if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0: if iteration % self.logging_cfg.checkpoint_frequency == 0: self._save_checkpoint(iteration) if getattr(self.cfg.experiment, "debug_sanity", False): self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") break if self.logging_cfg.save_model: model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model" self._save_model(model_path=model_path) self._close()