#!/usr/bin/env python3
# Filename: PPO_RC.py
# Author: Andrew Carr (direct questions to andrewcarr319@gmail.com after June 2026)
"""
A class that implements PPO for a MEMS-based reservoir computer for continuous
reinforcement learning tasks.
MEMS dynamics simulation is done with a rk4 DDE solver implemented
with JAX for speed.
Objects are created with, optionally, a list of hyperparameters, vectorized
environment from Farama Gymnasium, a reward function, previously
trained models, singular or specific hyperparameter values.
"""
# See the PPO_RC_pseudo pseudocode an easier-to-read version of this implementation.
import os
OMP_NUM_THREADS=1
MKL_NUM_THREADS=1
OPENBLAS_NUM_THREADS=1
NUMEXPR_NUM_THREADS=1
import numpy as np
import gymnasium as gym
from dataclasses import dataclass, asdict, replace, field
from typing import Any, Mapping, Optional, List
import wandb
import time
import json
import os
import jax
import jax.numpy as jnp
import optax
from hardware_rc.reservoir import Reservoir
import tempfile
from datetime import datetime
import math
import random
NORM_PRESETS = {
"CartPole-v1": {
"factor": [4.8*2,3+4,2*0.418,2+4],
"offset": [4.8,4,0.5,4]
},
"CartPole-v0": {
"factor": [4.8*2,3+4,2*0.418,2+4],
"offset": [4.8,4,0.5,4]
},
"LunarLander-v3": {
"factor": [ 2*2.5, 2*2.5, 2*10.0, 2*10.0, 2*6.2831855, 2*10.0, 2*1.0, 2*1.0 ],
"offset": [ 2.5, 2.5, 10.0, 10.0, 6.2831855, 10.0, 1.0, 1.0 ]
},
"MountainCar-v0": {
"factor": [2*1.2, 2*0.07],
"offset": [1.2, 0.07]
},
"MountainCarContinuous-v0": {
"factor": [2*1.2, 2*0.07],
"offset": [1.2, 0.07]
},
"BipedalWalker-v3": {
"factor": [2*3.1415927, 2*5., 2*5., 2*5., 2*3.1415927, 2*5., 2*3.1415927, 2*5., 2*5., 2*3.1415927, 2*5., 2*3.1415927, 2*5., 2*5., 2*1., 2*1., 2*1., 2*1., 2*1., 2*1., 2*1., 2*1., 2*1., 2*1.],
"offset": 24*[0],
}
}
__all__ = ["PPO_RC"]
@dataclass(frozen=True)
class PPOConfig:
# ---- model/solver ---
# DEFAULT HYPERPARAMETERS
env_name: str = "default_env"
general_seed: int = 1
mask_seed: int = 1
rc_seed: int = 1
env_seed: int = 1
val_size: int = 5
N_envs: int = 10
T_horizon: int = 2048
gamma: float = 0.99
lamb: float = 0.92
clip_eps: float = 0.15
v_clip_eps: float = 0.2
value_coef: float = 0.5
entropy_coef: float = 0.008
entropy_coef_step: float = 7.5e5
entropy_coef_min: float = 0.003
max_grad_norm: float = 2
num_epochs: int = 5
minibatch_size: int = 128
learning_rate: float = .0003
learning_rate_decay: float = 1
learning_rate_min: float = 1e-5
log_std_min: float = -4
log_std_max: float = 0
action_min: float = -1
action_max: float = 1
N: int = 300
theta: float = 1
T: float = 6*np.pi
h: float = 0.02
SampleDelay: int = 3
InputConnectivity: float = 0.2
NormalizationFactor: List[float] = field(default_factory=list)
NormalizationOffset: List[float] = field(default_factory=list)
amplification: int = 100
tau: int = 0
fb_gain: float = 0.0
trials: int = 100
def __post_init__(self):
"""This runs right after the object is initialized."""
# We use object.__setattr__ because the dataclass is frozen=True
if not self.NormalizationFactor or not self.NormalizationOffset:
preset = NORM_PRESETS.get(self.env_name, NORM_PRESETS["CartPole-v1"])
object.__setattr__(self, "NormalizationFactor", preset["factor"])
object.__setattr__(self, "NormalizationOffset", preset["offset"])
def validate(self) -> None:
if self.N <= 0:
raise ValueError(f"N must be > 0: {self.N}")
if self.minibatch_size <= 0:
raise ValueError(f"batch_size must be > 0: {self.minibatch_size}")
if self.gamma < 0:
raise ValueError(f"gamma must be >= 0: {self.gamma}")
@classmethod
def from_dict(cls, d: Mapping[str, Any]) -> "PPOConfig":
# ignore unknown keys gracefully
known = {k: v for k, v in d.items() if k in cls.__dataclass_fields__}
cfg = cls(**known) # type: ignore[arg-type]
cfg.validate()
return cfg
def updated(self, **overrides: Any) -> "PPOConfig":
# ignore unknown override keys as well
known_overrides = {k: v for k, v in overrides.items()
if k in self.__dataclass_fields__}
if not known_overrides:
return self
cfg = replace(self, **known_overrides)
cfg.validate()
return cfg
[docs]
class PPO_RC:
def __init__(self, env=None, *, reward_function=None, env_name=None, watch=True, N_envs=10,
config: Optional[PPOConfig | Mapping[str, Any]] = None,
model: Optional[str] = None,
**overrides: Any) -> None:
"""
A class that implements a PPO Reservoir Computer for continuous
reinforcement learning tasks.
Args:
env (gym.Env): The vectorized environment to train the model on.
reward_function (Optional[Callable]): Custom reward function.
config (Optional[PPOConfig | Mapping[str, Any]]): Hyperparameter configuration for the PPO RC.
model (Optional[str]): Path to a pre-trained model to load.
overrides: Any hyperparameter overrides.
"""
if env is None:
mode = None
if watch:
mode = "human"
if env_name is not None:
env = gym.make_vec(env_name, num_envs=N_envs, vectorization_mode='async', render_mode=mode)
else:
print("Warning: No environment provided. Defaulting to CartPole-v1.")
env = gym.make_vec("MountainCarContinuous-v0", num_envs=N_envs, vectorization_mode='async', render_mode=mode)
try:
self.env_name = env.unwrapped.spec.id
except:
self.env_name = env_name if env_name is not None else "unknown_env"
print(f"Warning: Could not determine environment name from env. Using '{self.env_name}' as env_name.")
# start from defaults
base = PPOConfig(env_name=self.env_name)
# merge user config (dataclass or dict)
if isinstance(config, PPOConfig):
base = config.updated(env_name=self.env_name)
elif isinstance(config, Mapping):
config["env_name"] = self.env_name
base = PPOConfig.from_dict(config)
elif config is not None:
raise TypeError("config must be a PPOConfig, dict-like, or None")
# apply overrides (highest priority)
print(f"Using env_name: {self.env_name}")
self.config = base.updated(**overrides)
# ----------------------------------------------------------------------------------------
### RL parameters
if self.config.general_seed != 0:
self.rc_seed = self.config.general_seed
self.mask_seed = self.config.general_seed
self.env_seed = self.config.general_seed
self.weight_seed = self.config.general_seed
else:
self.rc_seed = self.config.rc_seed
self.mask_seed = self.config.mask_seed
self.env_seed = self.config.env_seed
self.weight_seed = self.config.weight_seed
self.rand = np.random.RandomState(self.config.rc_seed)
self.env = env
self.test_env = None
self.trials = self.config.trials
self.state_shape = env.observation_space.shape[-1]
self.action_shape = env.action_space.shape[-1] # was .n before
self.N_envs = self.config.N_envs
self.T_horizon = self.config.T_horizon
self.gamma = self.config.gamma
self.lamb = self.config.lamb
self.clip_eps = self.config.clip_eps
self.v_clip_eps = self.config.v_clip_eps
self.value_coef = self.config.value_coef
self.entropy_coef = self.config.entropy_coef
self.entropy_coef_step = self.config.entropy_coef_step
self.entropy_coef_min = self.config.entropy_coef_min
self.entropy_coef_start = self.entropy_coef
self.max_grad_norm = self.config.max_grad_norm
self.learning_rate = self.config.learning_rate
self.learning_rate_decay = self.config.learning_rate_decay
self.learning_rate_min = self.config.learning_rate_min
self.action_min = self.config.action_min
self.action_max = self.config.action_max
self.num_epochs = self.config.num_epochs
self.minibatch_size = self.config.minibatch_size
self.timesteps = 0
self.val_size = self.config.val_size
self.trial_loss = 0
self.last_loss = 0
if reward_function is None:
def r_func(*, rewards, cur_state, new_state, actions, dones, hparams, validate=False):
return rewards
self.reward_function = r_func
else:
self.reward_function = reward_function
self.best_reward = float('-inf')
self.log_std_min = self.config.log_std_min
self.log_std_max = self.config.log_std_max
# ----------------------------------------------------------------------------------------
### MEMS RC parameters
self.N = self.config.N # number of neurons in reservoir
self.theta = self.config.theta # length of one sample of reservoir state [non-dimensional]
self.T = self.config.T # total simulation time for one reservoir state [non-dimensional]
self.h = self.config.h # integration step size
self.tau = self.config.tau
self.fb_gain = self.config.fb_gain
self.T_final = self.theta * self.N # total simulation time for one reservoir state [non-dimensional]
self.SampleDelay = self.config.SampleDelay # how long to wait before starting to sample the reservoir state (in number of h steps)
self.MEMS_IC = np.array([0.0,0.0]) # MEMS initial conditions (position, velocity)
self.tau_steps = int(self.tau * jnp.floor(self.theta/self.h))
m = self.tau_steps
buf_len = max(m+2, 2)
self.pos_buf = jnp.full((buf_len,), 0.0, dtype=jnp.float32)
self.vel_buf = jnp.full((buf_len,), 0.0, dtype=jnp.float32)
self.buf_idx = jnp.int32(0)
self.buf_cnt = jnp.int32(0) # buffer count
self.time_array = np.array([0]) # start with initial time of zero
self.MEMS_state = np.array([[0.0, 0.0]]) # initialize the array of MEMS state array (stores the MEMS dynamical response)
self.NormalizationFactor = self.config.NormalizationFactor
self.NormalizationOffset = self.config.NormalizationOffset
self.amplification = self.config.amplification
# Create buffers
self._buffers()
# ----------------------------------------------------------------------------------------
# check if model is provided after defining all hyperparams
self.given_model_path = None
if model is not None:
print('transfering over model')
root, extension = os.path.splitext(model)
if extension.lower() != '.npz':
raise TypeError('File extension must be .npz')
self._load_reservoir(model)
self.given_model_path = model
known_overrides = {k: v for k, v in overrides.items()
if k in self.__dict__}
print(f'overrides: {known_overrides}')
self.__dict__.update(known_overrides)
return
# ----------------------------------------------------------------------------------------
# input weight matrix aka mask
self.reservoir = Reservoir(h=self.h,
theta=self.theta,
N=self.N,
tau=self.tau,
fb_gain=self.fb_gain,
sd=self.SampleDelay,
amp=self.amplification,
norm_factor=self.NormalizationFactor,
norm_offset=self.NormalizationOffset,
state_shape=self.state_shape,
input_connectivity=self.config.InputConnectivity,
mask=True,
mask_seed=self.config.mask_seed,
normalize_mask=False,
)
# ----------------------------------------------------------------------------------------
# actor readout creation
scale = 1/np.sqrt(self.N)
self.Wa_mu = self.rand.normal(0, scale, size=(self.N, self.action_shape)) # shape = (N, DoF)
self.b_mu = np.zeros((self.action_shape,), dtype=np.float32)
self.Wa_std = self.rand.normal(0, 0.1*scale, size=(self.N, self.action_shape))
self.b_std = np.full((self.action_shape,), -0.5, dtype=np.float32)
# critic readout creation
self.Wc = self.rand.normal(0, scale, size=(self.N, 1)) # shape = (N, 1)
self.b_c = np.float32(0.0)
self.opt = optax.chain(
optax.clip_by_global_norm(self.max_grad_norm),
optax.adam(self.learning_rate)
)
self.params = self._get_params()
self.opt_state = self.opt.init(self.params)
def _get_params(self):
return {
"Wa_mu": self.Wa_mu,
"b_mu": self.b_mu,
"Wa_std": self.Wa_std,
"b_std": self.b_std,
"Wc": self.Wc,
"b_c": self.b_c,
}
def _set_params(self, params):
self.Wa_mu = params["Wa_mu"]
self.b_mu = params["b_mu"]
self.Wa_std = params["Wa_std"]
self.b_std = params["b_std"]
self.Wc = params["Wc"]
self.b_c = params["b_c"]
def _buffers(self):
"""Create the buffers for training."""
self.S_buf = np.zeros((self.T_horizon, self.N_envs, self.state_shape))
self.A_buf = np.zeros((self.T_horizon, self.N_envs, self.action_shape))
self.R_buf = np.zeros((self.T_horizon, self.N_envs))
self.DONE_buf = np.zeros((self.T_horizon, self.N_envs))
self.LOGP_buf = np.zeros((self.T_horizon, self.N_envs))
self.V_buf = np.zeros((self.T_horizon, self.N_envs))
self.ADV_buf = np.zeros((self.T_horizon, self.N_envs))
self.RET_buf = np.zeros((self.T_horizon, self.N_envs)) # for critic training
self.X_buf = np.zeros((self.T_horizon, self.N_envs, self.N))
def _flatten_buffers(self):
"""Flatten the buffers for training."""
self.S_buf = self.S_buf.reshape(-1, self.state_shape)
self.A_buf = self.A_buf.reshape(-1, self.action_shape)
self.LOGP_buf = self.LOGP_buf.reshape(-1)
self.ADV_buf = self.ADV_buf.reshape(-1)
self.RET_buf = self.RET_buf.reshape(-1)
self.V_buf = self.V_buf.reshape(-1)
self.X_buf = self.X_buf.reshape(-1, self.N)
def _calc_pi(self, x):
"""Calculate mean and log standard deviation of the current policy (Gaussian)."""
mu = x @ self.Wa_mu + self.b_mu # returns each with shape (1,action_shape)
log_std_raw = x @ self.Wa_std + self.b_std # returns each with shape (1,action_shape)
log_std = np.clip(log_std_raw, self.log_std_min, self.log_std_max)
std = np.exp(log_std)
return mu, log_std, std
def _calc_log_prob(self, mu, log_std, std, a):
"""Calculate log probability of action a under the current policy (Gaussian)."""
return -0.5 * np.sum(np.square(a - mu) / np.square(std) + 2*log_std + np.log(2*np.pi))
def _rollout(self):
"""Collect trajectories by running the current policy in the environment."""
self._buffers()
a = np.zeros((self.N_envs, self.action_shape))
for t in range(self.T_horizon):
for N_env in range(self.N_envs):
x = self.reservoir.sim(obs=self.obs[N_env], mode='train', env_idx=N_env) # neuron outputs for this env
mu, log_std, std = self._calc_pi(x)
a_exact = mu + std * self.rand.normal(size=mu.shape) # get a random action within a std
a[N_env] = a_exact # np.clip(a_exact, self.action_min, self.action_max)
# self.debug_check(
# label=f"rollout t={t} env={N_env}",
# obs=self.obs[N_env],
# x=x,
# mu=mu,
# log_std=log_std,
# std=std,
# action=a_exact,
# raise_on_error=False,
# )
# calc log probability of a under current policy
logp = self._calc_log_prob(mu, log_std, std, a_exact) ## changed from a[N_env]
v = (x @ self.Wc).squeeze(-1) + self.b_c # calc value function
self.X_buf[t, N_env] = x
self.S_buf[t, N_env] = self.obs[N_env]
self.A_buf[t, N_env] = a[N_env]
self.LOGP_buf[t, N_env] = logp
self.V_buf[t, N_env] = v
last_obs = self.obs
self.obs, rewards, terminations, truncations, _ = self.v_env.step(a)
dones = terminations | truncations
self.reservoir._zero_reservoir(mode='train', env_idx=dones)
rewards = self.reward_function(rewards=rewards, cur_state=last_obs, new_state=self.obs, actions=self.A_buf[t], dones=dones, hparams=self.config)
self.R_buf[t] = rewards
self.DONE_buf[t] = dones
# sys.exit()
# bootstap value function for last value in T_horizon
v_last_buf = np.zeros((self.N_envs))
for N_env in range(self.N_envs):
v_last_buf[N_env] = (self.reservoir.sim(obs=self.obs[N_env], mode='train', env_idx=N_env, update_state=False) @ self.Wc).squeeze(-1) + self.b_c
self.timesteps += self.N_envs * self.T_horizon
return v_last_buf
def _calc_adv_ret(self, v_last):
"""Calculate advantages and returns using Generalized Advantage Estimation (GAE)."""
gae = np.zeros((self.N_envs))
# A_t = del_t + (gamma*lambda)*del_t+1 + ... + (gamma*lambda)^(T_horizon-t+1)*del_T-1
# del_t = r_t + gamma*V(s_t+1) - V(s_t)
# To do this algorithmically, we start at the end of T_horizon because at t=0,
# we need to know all the advantages to come. Starting at the end is the most efficient
# way to do this computation
# At each time step, we are comparing the actor to the critic, so even though a timestep
# further in the future will have less advantage, the critic will do compute
# something in the same domain so the comparison is valid and useful.
# critic defines the expectation, actor learns to beat that expectation
for t in reversed(range(self.T_horizon)):
nonterminal = 1 - self.DONE_buf[t]
v_next = v_last if t == self.T_horizon - 1 else self.V_buf[t+1]
del_t = self.R_buf[t] + self.gamma * v_next * nonterminal - self.V_buf[t]
gae = del_t + self.gamma * self.lamb * nonterminal * gae
self.ADV_buf[t] = gae
self.RET_buf = self.ADV_buf + self.V_buf
# Normalize advantages and returns per batch to mean=0 and unit variance
# Less bias in training and we get a relative ranking of actions
self.ADV_buf = (self.ADV_buf - np.mean(self.ADV_buf)) / (np.std(self.ADV_buf) + 1e-8)
# self.RET_buf = (self.RET_buf - np.mean(self.RET_buf)) / (np.std(self.RET_buf) + 1e-8)
# ------------------ Auto Diff w/ Jax specifics -----------------------------
@staticmethod
def _forward(params, x, log_std_min, log_std_max):
"""'Forward pass' of the RC policy network."""
mu = x @ params['Wa_mu'] + params['b_mu']
log_std_b = x @ params['Wa_std'] + params['b_std']
log_std_b = jnp.clip(log_std_b, log_std_min, log_std_max)
std = jnp.exp(log_std_b)
v = (x @ params["Wc"]).squeeze(-1) + params['b_c']
return mu, log_std_b, std, v
@staticmethod
def _log_prob_gauss(a, mu, log_std_b, std):
"""Calculate the log probability of action a under a Gaussian distribution with mean mu and standard deviation std."""
return -0.5 * jnp.sum(jnp.square(a - mu) / jnp.square(std) + 2*log_std_b + jnp.log(2*jnp.pi), axis=-1)
@staticmethod
def _entropy_gauss(log_std_b):
"""Calculate the entropy of a Gaussian distribution given its log standard deviation."""
return (0.5 * (1.0 + jnp.log(2.0*jnp.pi)) + log_std_b).sum(-1)
@staticmethod
def _ppo_loss(params, x_mb, a_mb, logp_mb, adv_mb, ret_mb,
clip_eps, v_clip_eps, value_coef, entropy_coef, v_old_mb,
log_std_max, log_std_min):
"""
Calculate the PPO Loss for a given minibatch of data and model parameters.
The PPO loss consists of three components:
1. Policy Loss: Encourages the new policy to improve upon the old policy while preventing large
updates that could destabilize training. This is done using a clipped surrogate objective.
2. Value Loss: Encourages the new value function to better represent the expected returns from the states.
This is typically a mean squared error loss between the predicted values and the normalized returns,
with optional clipping to prevent large updates.
3. Entropy Loss: Encourages exploration by maximizing the entropy of the policy's action distribution.
This helps prevent premature convergence to suboptimal policies.
"""
# calc actor and critic outputs
mu, log_std_b, std, v_new = PPO_RC._forward(params, x_mb, log_std_min, log_std_max)
# calc log prob of action taken from rollout
logp_new = PPO_RC._log_prob_gauss(a_mb, mu, log_std_b, std)
# calc probability ratio
ratio = jnp.exp(logp_new - logp_mb) # r_t(θ) = πnew/πold
# clipped surrogate objective
# objective = loss function = ratio*Adv
# we use the expectation of this (average over mb)
obj_unclipped = ratio*adv_mb
# keep ratio within 1-eps to 1+eps
obj_clipped = jnp.clip(ratio, 1-clip_eps, 1+clip_eps)*adv_mb
# clipping only applies if obj_clipped is less than obj_unclipped
policy_loss = -jnp.mean(jnp.minimum(obj_unclipped, obj_clipped))
# normalize returns and values, (ret is the baseline)
ret_mean = jnp.mean(ret_mb)
ret_std = jnp.std(ret_mb) + 1e-8
ret_n = (ret_mb - ret_mean) / ret_std
v_new_n = (v_new - ret_mean) / ret_std
v_old_n = (v_old_mb - ret_mean) / ret_std
v_pred = v_new_n
v_pred_clipped = v_old_n + jnp.clip(v_pred - v_old_n, -v_clip_eps, v_clip_eps)
v_err = (ret_n - v_pred)**2
v_err_clipped = (ret_n - v_pred_clipped)**2
value_loss = 0.5 * value_coef * jnp.mean(jnp.maximum(v_err, v_err_clipped))
# value_loss = 0.5*value_coef*jnp.mean((ret_mb - v_new)**2)
# entropy_loss = -entropy_coef*jnp.mean(0.5*2*log_std_b + jnp.log(2*jnp.pi))
entropy = PPO_RC._entropy_gauss(log_std_b) # (batch,)
entropy_loss = -entropy_coef * jnp.mean(entropy)
# print(f'Policy Loss: {policy_loss[0]}, Value Loss: {value_loss[0]}, Entropy Loss: {entropy_loss[0]}')
return policy_loss + value_loss + entropy_loss, policy_loss, value_loss, entropy_loss
def _adam_step(self, x_mb, a_mb, logp_mb, adv_mb, ret_mb, v_old_mb):
def _loss_only(g_params):
return PPO_RC._ppo_loss(g_params, x_mb, a_mb, logp_mb, adv_mb,
ret_mb, self.clip_eps, self.v_clip_eps,
self.value_coef, self.entropy_coef, v_old_mb,
self.log_std_max, self.log_std_min)[0]
# create pytree for jax to take gradients with respect to
g_params = {'Wa_mu': self.Wa_mu,
'b_mu': self.b_mu,
'Wa_std': self.Wa_std,
'b_std': self.b_std,
'Wc': self.Wc,
'b_c': self.b_c}
_, p_loss, v_loss, e_loss = PPO_RC._ppo_loss(g_params, x_mb, a_mb, logp_mb, adv_mb,
ret_mb, self.clip_eps, self.v_clip_eps,
self.value_coef, self.entropy_coef, v_old_mb,
self.log_std_max, self.log_std_min)
# print(f'policy_loss: {p_loss}, value_loss: {v_loss}, entropy_loss: {e_loss}')
loss, grads = jax.value_and_grad(_loss_only)(g_params)
updates, self.opt_state = self.opt.update(grads, self.opt_state, self.params)
self.params = optax.apply_updates(self.params, updates)
self._set_params(self.params)
return loss
def _optimize(self):
# get random permutation of values from 0 to len(A_buf) for training
idx = self.rand.permutation(len(self.A_buf))
batch_loss = [] # could change this to a fixed length if slow
batch_scales = []
for epoch in range(self.num_epochs):
for batch_start in range(0, len(self.A_buf), self.minibatch_size):
mb = idx[batch_start:batch_start+self.minibatch_size]
s_mb = jnp.asarray(self.S_buf[mb])
a_mb = jnp.asarray(self.A_buf[mb])
logp_mb = jnp.asarray(self.LOGP_buf[mb])
adv_mb = jnp.asarray(self.ADV_buf[mb])
ret_mb = jnp.asarray(self.RET_buf[mb])
v_old_mb = jnp.asarray(self.V_buf[mb], jnp.float32)
x_mb = jnp.asarray(self.X_buf[mb], jnp.float32)
# get reservoir outputs
mb_size = len(mb)
# x_mb = np.zeros((mb_size, self.N))
# i=0
# for state in s_mb:
# x_mb[i] = self.reservoir.sim(state)
# i+=1
# print(f's_mb: {s_mb.shape}')
# print(f'x_mb: {x_mb.shape}')
# print(f'a_mb: {a_mb.shape}')
# print(f'batch_start: {batch_start}')
# reverse autodiff using jax
loss = self._adam_step(x_mb, a_mb, logp_mb, adv_mb, ret_mb, v_old_mb)
batch_loss.append(loss)
# batch_scales.append(scale)
# sys.exit()
trial_loss = sum(batch_loss)/len(batch_loss)
# trial_scale = sum(batch_scales)/len(batch_scales)
last_loss = loss
return trial_loss, last_loss, None # trial_scale
def _decay(self):
# Learning rate decay
self.learning_rate = max(self.learning_rate_min, self.learning_rate * self.learning_rate_decay)
# Entropy coefficient decay
k = math.log(2) / self.entropy_coef_step
self.entropy_coef = max(self.entropy_coef_min, self.entropy_coef_start * math.exp(-k * self.timesteps))
[docs]
def train(self, *, hparams=None, v_env=None, meta={}, trials=10, T_horizon=None,
num_epochs=None, folder_path=None, wandb_on=True, save_model=True,
ray_on=False, ray_metric=None, ray_mode='max',
save_best_reward=True, ray_report_freq=1,
resume=False, wandb_id=None):
"""
Train the agent using PPO with the reservoir computer.
Args:
hparams (Optional[Mapping[str, Any]]): Hyperparameters for training.
v_env (Optional[gym.Env]): A separate vectorized environment for training. If None, the current environment is used.
meta (dict): Metadata for logging and saving.
trials (int): Number of training trials. Defaults to agent's config.
T_horizon (Optional[int]): Maximum number of timesteps per episode. Defaults to agent's config.
num_epochs (Optional[int]): Number of epochs per training trial. Defaults to agent's config.
folder_path (Optional[str]): Path to save models. Defaults to 'PPO_RC/models/{env_name}-{group}'.
wandb_on (bool): Whether to use wandb for logging.
save_model (bool): Whether to save the model after training.
resume (bool): Whether to resume training from a checkpoint.
wandb_id (str): The wandb run id to resume training from.
Returns:
best_model_path (str): The path to the best model found during training.
"""
if ray_on:
try:
from ray.air import session
from ray.train import Checkpoint
from ray.air.integrations.wandb import WandbLoggerCallback
except:
print("Ray is not installed. Please install ray to use ray features.")
session = None
Checkpoint = None
WandbLoggerCallback = None
# wandb = None
ray_on = False
if v_env is not None:
self.v_env = v_env
self.N_envs = v_env.num_envs
else:
self.v_env = self.env
self.N_envs = self.v_env.num_envs
self.reservoir._init_reservoir_state(self.N_envs)
if T_horizon is not None:
self.T_horizon = T_horizon
if num_epochs is not None:
self.num_epochs = num_epochs
required_fields = ["project", "group", "job_type", "run_name", "tags"]
for field in required_fields:
if field not in meta:
meta[field] = "not_provided"
if folder_path is not None:
self.folder_path = folder_path
elif save_model:
print("No folder provided. Model will be saved to 'PPO_RC/models' in current directory.")
self.folder_path = f"PPO_RC/models/{self.env_name}-{meta['group']}"
if ray_on and ray_metric is None:
ray_metric = 'avg_reward'
if hparams is None:
hparams = self.config
datetime_str = datetime.now().strftime('%Y%m%d-%H%M%S')
run_animal = self._get_animal()
if not resume:
meta["run_name"] = f"{meta['run_name']}-{run_animal}-{datetime_str}"
# setup wandb
if wandb_on:
run = wandb.init(
project=meta["project"],
group=meta["group"] if not resume else None,
job_type=meta["job_type"] if not resume else None,
name=meta["run_name"] if not resume else None,
tags=meta['tags'] if not resume else None,
config={"hp": hparams, "meta": meta} if not resume else None,
id=wandb_id if resume else None,
resume= "must" if resume else False,
)
self.run_name = meta["run_name"] if not resume else f'{run.name}_resume'
eval_metric = float('-inf')
if ray_on and ray_mode == 'min':
eval_metric = float('inf')
reward_metric = float('-inf')
self.obs, _ = self.v_env.reset(seed=self.config.env_seed)
for trial in range(trials):
trial_time = time.time_ns()
# rollout envs for training data
# print('rolling out')
rollout_time = time.time_ns()
v_last = self._rollout()
rollout_time = (time.time_ns() - rollout_time)/1e6 # ms
# calculate advantages and returns
# print('calc_adv_ret')
adv_ret_time = time.time_ns()
self._calc_adv_ret(v_last)
adv_ret_time = (time.time_ns() - adv_ret_time)/1e6 # ms
# Flatten S_buf, A_buf, LOGP_OLD_buf, ADV_buf and RET_buf for training
# print('flatten')
flatten_time = time.time_ns()
self._flatten_buffers()
flatten_time = (time.time_ns() - flatten_time)/1e6 # ms
# optimize loss function
# print('opt')
opt_time = time.time_ns()
trial_loss, last_loss, trial_scale = self._optimize()
opt_time = (time.time_ns() - opt_time)/1e9 # s
# validate
# print('validate')
val_time = time.time_ns()
val_results = self._validate()
val_time = (time.time_ns() - val_time)/1e9 # s
# log results
if val_results['avg_reward'] >= reward_metric and not ray_on and save_best_reward:
reward_metric = val_results['avg_reward']
self.best_reward = val_results['avg_reward']
if resume:
name_val = f"resume_best_{datetime_str}"
else:
name_val = f"best_{run_animal}_{datetime_str}"
if save_model:
self._save_reservoir(name_val=name_val)
best_model_path = self.model_path
print('new best avg reward model:')
# Save last model always
if resume and not ray_on:
name_val = f"resume_last_{datetime_str}"
else:
name_val = f"last_{run_animal}_{datetime_str}"
if save_model and not ray_on:
self._save_reservoir(name_val=name_val)
avg_action_1 = np.mean(val_results['actual_actions'][:,0])
action_std_1 = np.std(val_results['actual_actions'][:,0])
if self.action_shape >=2:
avg_action_2 = np.mean(val_results['actual_actions'][:,1])
action_std_2 = np.std(val_results['actual_actions'][:,1])
if self.action_shape >= 3:
avg_action_3 = np.mean(val_results['actual_actions'][:,2])
action_std_3 = np.std(val_results['actual_actions'][:,2])
if self.action_shape == 4:
avg_action_4 = np.mean(val_results['actual_actions'][:,3])
action_std_4 = np.std(val_results['actual_actions'][:,3])
if wandb_on:
wandb.log({
'avg_reward': val_results['avg_reward'],
'avg_trial_length': val_results['avg_trial_length'],
'median_trial_length': val_results['median_trial_length'],
'neuron_sat': val_results['avg_neuron_sat'],
'avg_log_std': val_results['avg_log_std'],
'trial_loss': trial_loss,
'last_loss': last_loss,
'trial': trial,
'avg_action': val_results['avg_action'],
'action_std': val_results['action_std'],
'action_min': val_results['action_min'],
'action_max': val_results['action_max'],
'avg_value': val_results['avg_value'],
'trial_scale': trial_scale,
'eval_metric': eval_metric,
'entropy_coef': self.entropy_coef,
'avg_action_1': avg_action_1,
'avg_action_2': avg_action_2 if self.action_shape >= 2 else None,
'avg_action_3': avg_action_3 if self.action_shape >= 3 else None,
'avg_action_4': avg_action_4 if self.action_shape >= 4 else None,
'action_std_1': action_std_1,
'action_std_2': action_std_2 if self.action_shape >= 2 else None,
'action_std_3': action_std_3 if self.action_shape >= 3 else None,
'action_std_4': action_std_4 if self.action_shape >= 4 else None,
'rollout_time': rollout_time,
'adv_ret_time': adv_ret_time,
'flatten_time': flatten_time,
'opt_time': opt_time,
'val_time': val_time,
})
if ray_on:
# if ray_mode is 'max', we want to maximize eval_metric, vice versa for 'min'
if val_results[ray_metric] >= eval_metric and ray_mode == 'max' or val_results[ray_metric] <= eval_metric and ray_mode == 'min':
eval_metric = val_results[ray_metric]
with tempfile.TemporaryDirectory() as tmpdir:
# 2) save your artifact(s) into that folder
filename = f"reservoir_best-{meta['run_name']}.npz"
out_path = os.path.join(tmpdir, filename)
self._save_reservoir(file_path=out_path)
session.report({
'avg_reward': val_results['avg_reward'],
'avg_trial_length': val_results['avg_trial_length'],
'neuron_sat': val_results['avg_neuron_sat'],
'avg_log_std': val_results['avg_log_std'],
'trial_loss': trial_loss,
'last_loss': last_loss,
'eval_metric_cur': val_results[ray_metric],
'trial': trial,
'avg_action': val_results['avg_action'],
'action_std': val_results['action_std'],
'avg_value': val_results['avg_value'],
'trial_scale': trial_scale,
'eval_metric': eval_metric,
}, checkpoint=Checkpoint.from_directory(tmpdir))
elif trial % ray_report_freq == 0:
session.report({
'avg_reward': val_results['avg_reward'],
'avg_trial_length': val_results['avg_trial_length'],
'neuron_sat': val_results['avg_neuron_sat'],
'avg_log_std': val_results['avg_log_std'],
'trial_loss': trial_loss,
'last_loss': last_loss,
'eval_metric_cur': val_results[ray_metric],
'trial': trial,
'avg_action': val_results['avg_action'],
'action_std': val_results['action_std'],
'avg_value': val_results['avg_value'],
'trial_scale': trial_scale,
'eval_metric': eval_metric,
})
# decay hyperparameters if desired
self._decay()
trial_time = (time.time_ns() - trial_time)/1e9 # s
if not ray_on:
print(f"trial: {trial+1}, avg_reward: {val_results['avg_reward']:.2f}, avg_trial_length: {val_results['avg_trial_length']:.0f}, neuron_sat: {val_results['avg_neuron_sat']:.2f}, action u+/-sig: {val_results['avg_action']:.2f} +/- {val_results['action_std']:.3f}")
print(f'Trial Time: {trial_time:.1f} [s], rollout time: {rollout_time/1000:.1f} [s], adv_ret time: {adv_ret_time:.2f} [ms], opt time: {opt_time:.1f} [s], val time: {val_time:.1f} [s]')
if save_model and not ray_on:
return best_model_path
elif not save_model or not ray_on:
print('Model saving disabled')
return None
else:
return None
def _act(self, obs, mode, env_idx=0):
x = self.reservoir.sim(obs=obs, mode=mode, env_idx=env_idx)
mu, log_std, _ = self._calc_pi(x)
v = (x @ self.Wc).squeeze(-1) + self.b_c
return mu, log_std, x, v
def _validate(self):
if self.test_env is None:
if self.v_env.unwrapped.spec.id == 'LunarLander-v3':
self.test_env = gym.make(f'{self.v_env.unwrapped.spec.id}', continuous=True)
else:
self.test_env = gym.make(f'{self.v_env.unwrapped.spec.id}')
avg_neuron_sats = []
avg_log_stds = []
rewards = []
trial_lengths = []
avg_actions = []
avg_values = []
avg_action_std = []
action_maxs = []
action_mins = []
start = 20
end = start + self.val_size
rang = range(start,end,1)
for i in rang:
done = False
env_seed = i
tot_reward = 0
trial_length = 0
tot_neuron_sat = 0
tot_log_std = 0
action_sum = 0
value_sum = 0
actual_actions = []
obs, _ = self.test_env.reset(seed=env_seed)
self.reservoir._zero_reservoir(mode='val') # reset reservoir so deterministic per seed
while not done:
# print(obs)
a, log_std, x, v = self._act(obs, mode='val')
# self.debug_check(
# label=f"validate seed={env_seed} step={trial_length}",
# obs=obs,
# x=x,
# mu=a,
# log_std=log_std,
# std=np.exp(log_std),
# action=a,
# raise_on_error=False,
# )
last_obs = obs
obs, reward, termination, truncation, _ = self.test_env.step(a)
done = termination or truncation
# print(reward)
reward = self.reward_function(rewards=reward, cur_state=last_obs, new_state=obs, actions=a, dones=done, hparams=self.config, validate=True)
tot_reward += reward
trial_length += 1
tot_neuron_sat += np.sum(np.abs(x)>=0.89)*100/x.size
tot_log_std += float(np.mean(log_std))
action_sum += np.sum(a)
value_sum += v
actual_actions.append(a)
rewards.append(tot_reward)
avg_neuron_sats.append(tot_neuron_sat/trial_length)
avg_log_stds.append(tot_log_std/trial_length)
trial_lengths.append(trial_length)
avg_actions.append(action_sum/trial_length)
avg_values.append(value_sum/trial_length)
avg_action_std.append(np.std(actual_actions))
action_mins.append(np.min(actual_actions))
action_maxs.append(np.max(actual_actions))
# print(f'Validate trial {i-min+1}: trial length: {trial_length}, reward: {tot_reward:.3f}, neuron sat: {tot_neuron_sat/trial_length:.2f}, avg_log_std: {tot_log_std/trial_length}')
val_results = {
'avg_reward': sum(rewards)/len(rewards),
'avg_neuron_sat': sum(avg_neuron_sats)/len(avg_neuron_sats),
'avg_log_std': sum(avg_log_stds)/len(avg_log_stds),
'avg_trial_length': sum(trial_lengths)/len(trial_lengths),
'trial_lengths': trial_lengths,
'median_trial_length': np.median(trial_lengths),
'avg_action': sum(avg_actions)/len(avg_actions),
'actual_actions': np.asarray(actual_actions),
'action_std': np.mean(avg_action_std),
'action_min': np.min(action_mins),
'action_max': np.max(action_maxs),
'avg_value': sum(avg_values)/len(avg_values),
}
return val_results
def _save_reservoir(self, *, name_val=None, file_path=None):
self.hparams = {
'algorithm': 'PPO_RC',
'environment': self.env.unwrapped.spec.id,
'run_name': self.run_name,
}
self.hparams.update(asdict(self.config))
meta = json.dumps(self.hparams)
if file_path is not None:
self.model_path = file_path
else:
file_name = f'{name_val}.npz' if name_val is not None else f'{self.run_name}.npz'
save_folder = f'{self.folder_path}/{self.run_name}'
os.makedirs(save_folder, exist_ok=True)
self.model_path = f'{save_folder}/{file_name}'
np.savez_compressed(self.model_path,
Wa_mu=self.Wa_mu,
b_mu=self.b_mu,
Wa_std=self.Wa_std,
b_std=self.b_std,
Wc=self.Wc,
b_c=self.b_c,
reservoir=self.reservoir,
meta=meta)
def _load_reservoir(self, path_npz, get_params=False):
with np.load(path_npz, allow_pickle=True) as data:
self.reservoir = data['reservoir'].item() if 'reservoir' in data else None
mask = data['mask'] if 'mask' in data else None
self.Wa_mu = data['Wa_mu']
self.b_mu = data['b_mu'] if 'b_mu' in data else None
self.Wa_std = data['Wa_std']
self.b_std = data['b_std'] if 'b_std' in data else None
self.Wc = data['Wc']
self.b_c = data['b_c'] if 'b_c' in data else None
# Decode JSON string from the NPZ entry
meta_dict = json.loads(data['meta'].item())
# Set attributes dynamically
for key, value in meta_dict.items():
setattr(self, key, value)
if self.reservoir is None:
self.reservoir = Reservoir(h=self.h,
theta=self.theta,
N=self.N,
tau=self.tau,
fb_gain=self.fb_gain,
sd=self.SampleDelay,
amp=self.amplification,
norm_factor=self.NormalizationFactor,
norm_offset=self.NormalizationOffset,
state_shape=self.state_shape,
load=True,
mask=mask)
self.opt = optax.chain(
optax.clip_by_global_norm(self.max_grad_norm),
optax.adam(self.learning_rate)
)
self.params = self._get_params()
self.opt_state = self.opt.init(self.params)
# Store as a Python dict
self.hparams = meta_dict
if get_params:
return self
def _get_animal(self):
return random.choice(animals)
animals = ['Canidae', 'Felidae', 'Cat', 'Cattle', 'Dog', 'Donkey', 'Goat', 'Horse', 'Pig', 'Rabbit',
'Aardvark', 'Aardwolf', 'Albatross', 'Alligator', 'Alpaca', 'Amphibian', 'Anaconda',
'Angelfish', 'Anglerfish', 'Ant', 'Anteater', 'Antelope', 'Antlion', 'Ape', 'Aphid',
'Armadillo', 'Asp', 'Baboon', 'Badger', 'Bandicoot', 'Barnacle', 'Barracuda', 'Basilisk',
'Bass', 'Bat', 'Bear', 'Beaver', 'Bedbug', 'Bee', 'Beetle', 'Bird', 'Bison', 'Blackbird',
'Boa', 'Boar', 'Bobcat', 'Bobolink', 'Bonobo', 'Bovid', 'Bug', 'Butterfly', 'Buzzard',
'Camel', 'Canid', 'Capybara', 'Cardinal', 'Caribou', 'Carp', 'Cat', 'Catshark',
'Caterpillar', 'Catfish', 'Cattle', 'Centipede', 'Cephalopod', 'Chameleon', 'Cheetah',
'Chickadee', 'Chicken', 'Chimpanzee', 'Chinchilla', 'Chipmunk', 'Clam', 'Clownfish',
'Cobra', 'Cockroach', 'Cod', 'Condor', 'Constrictor', 'Coral', 'Cougar', 'Cow', 'Coyote',
'Crab', 'Crane', 'Crawdad', 'Crayfish', 'Cricket', 'Crocodile', 'Crow', 'Cuckoo', 'Cicada',
'Damselfly', 'Deer', 'Dingo', 'Dinosaur', 'Dog', 'Dolphin', 'Donkey', 'Dormouse', 'Dove',
'Dragonfly', 'Dragon', 'Duck', 'Eagle', 'Earthworm', 'Earwig', 'Echidna', 'Eel', 'Egret',
'Elephant', 'Elk', 'Emu', 'Ermine', 'Falcon', 'Ferret', 'Finch', 'Firefly', 'Fish',
'Flamingo', 'Flea', 'Fly', 'Flyingfish', 'Fowl', 'Fox', 'Frog', 'Gamefowl', 'Galliform',
'Gazelle', 'Gecko', 'Gerbil', 'Gibbon', 'Giraffe', 'Goat', 'Goldfish', 'Goose', 'Gopher',
'Gorilla', 'Grasshopper', 'Grouse', 'Guan', 'Guanaco', 'Guineafowl', 'Gull', 'Guppy',
'Haddock', 'Halibut', 'Hamster', 'Hare', 'Harrier', 'Hawk', 'Hedgehog', 'Heron', 'Herring',
'Hippopotamus', 'Hookworm', 'Hornet', 'Horse', 'Hoverfly', 'Hummingbird', 'Hyena', 'Iguana',
'Impala', 'Jackal', 'Jaguar', 'Jay', 'Jellyfish', 'Junglefowl', 'Kangaroo', 'Kingfisher',
'Kite', 'Kiwi', 'Koala', 'Koi', 'Krill', 'Ladybug', 'Lamprey', 'Landfowl', 'Lark', 'Leech',
'Lemming', 'Lemur', 'Leopard', 'Leopon', 'Limpet', 'Lion', 'Lizard', 'Llama', 'Lobster',
'Locust', 'Loon', 'Louse', 'Lungfish', 'Lynx', 'Macaw', 'Mackerel', 'Magpie', 'Mammal',
'Manatee', 'Mandrill', 'Marlin', 'Marmoset', 'Marmot', 'Marsupial', 'Marten', 'Mastodon',
'Meadowlark', 'Meerkat', 'Mink', 'Minnow', 'Mite', 'Mockingbird', 'Mole', 'Mollusk',
'Mongoose', 'Monkey', 'Moose', 'Mosquito', 'Moth', 'Mouse', 'Mule', 'Muskox', 'Narwhal',
'Newt', 'Nightingale', 'Ocelot', 'Octopus', 'Opossum', 'Orangutan', 'Orca', 'Ostrich',
'Otter', 'Owl', 'Ox', 'Panda', 'Panther', 'Parakeet', 'Parrot', 'Parrotfish', 'Partridge',
'Peacock', 'Peafowl', 'Pelican', 'Penguin', 'Perch', 'Pheasant', 'Pig', 'Pigeon', 'Pike',
'Pinniped', 'Piranha', 'Planarian', 'Platypus', 'Pony', 'Porcupine', 'Porpoise', 'Possum',
'Prawn', 'Primate', 'Ptarmigan', 'Puffin', 'Puma', 'Python', 'Quail', 'Quelea', 'Quokka',
'Rabbit', 'Raccoon', 'Rat', 'Rattlesnake', 'Raven', 'Reindeer', 'Reptile', 'Rhinoceros',
'Roadrunner', 'Rodent', 'Rook', 'Rooster', 'Roundworm', 'Sailfish', 'Salamander', 'Salmon',
'Sawfish', 'Scallop', 'Scorpion', 'Seahorse', 'Shark', 'Sheep', 'Shrew', 'Shrimp',
'Silkworm', 'Silverfish', 'Skink', 'Skunk', 'Sloth', 'Slug', 'Smelt', 'Snail', 'Snake',
'Snipe', 'Sole', 'Sparrow', 'Spider', 'Spoonbill', 'Squid', 'Squirrel', 'Starfish',
'Stingray', 'Stoat', 'Stork', 'Sturgeon', 'Swallow', 'Swan', 'Swift', 'Swordfish',
'Swordtail', 'Tahr', 'Takin', 'Tapir', 'Tarantula', 'Tarsier', 'Termite', 'Tern', 'Thrush',
'Tick', 'Tiger', 'Tiglon', 'Toad', 'Tortoise', 'Toucan', 'Trout', 'Tuna', 'Turkey',
'Turtle', 'Tyrannosaurus', 'Urial', 'Vicuna', 'Viper', 'Vole', 'Vulture', 'Wallaby',
'Walrus', 'Wasp', 'Warbler', 'Weasel', 'Whale', 'Whippet', 'Whitefish', 'Wildcat',
'Wildebeest', 'Wildfowl', 'Wolf', 'Wolverine', 'Wombat', 'Woodpecker', 'Worm', 'Wren',
'Xerinae', 'Yak', 'Zebra', 'Alpaca', 'Cat', 'Cattle', 'Chicken', 'Dog', 'Donkey', 'Ferret',
'Gayal', 'Goldfish', 'Guppy', 'Horse', 'Koi', 'Llama', 'Sheep', 'Yak']