#!/usr/bin/env python3
# Filename: DQN_RC.py
# Author: Andrew Carr (direct questions to andrewcarr319@gmail.com after June 2026)
"""
A class that implements DQN for a MEMS-based reservoir computer for
discrete 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,
environment from Farama Gymnasium, previously trained models, and
singular or specific hyperparameter values.
"""
from datetime import datetime
import os
import random
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 collections import deque
import wandb
from dataclasses import dataclass, asdict, replace, field
from typing import Any, Mapping, Optional, List
import time
import json
import tempfile
import jax
import jax.numpy as jnp
from hardware_rc.reservoir import Reservoir
### Normalization factor and offset presets for different environments, can be expanded as needed
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]
},
}
__all__ = ["DQN_RC"]
@dataclass(frozen=True)
class DQNConfig:
# ---- model/solver ---
# DEFAULT HYPERPARAMETERS
env_name: str = "default_env"
rc_seed: int = 1
mask_seed: int = 1
weight_seed: int = 1
env_seed: int = 1
general_seed: int = 1
N: int = 300
bufferLength: int = 6000
learning_rate: float = 1e-4
learning_rate_decay: float = 1
learning_rate_min: float = 1e-5
epsilon: float = 1
epsilon_min: float = 0.01
epsilon_decay: float = 0.999
gamma: float = 0.99
beta: float = 100
vel_weight: float = .5
target_update_rate: int = 4
batch_size: int = 16
theta: float = 1
T: float = 6*np.pi
h: float = 0.02
SampleDelay: int = 0
InputConnectivity: float = 0.2
rewardNormalizationFactor: float = 1.0
NormalizationFactor: List[float] = field(default_factory=list)
NormalizationOffset: List[float] = field(default_factory=list)
amplification: float = 360.0
tau: int = 0
fb_gain: float = 0
tau_N: float = 0.0 # if non-zero, tau is set to tau_N * N
trials: int = 100
val_size: int = 10
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.batch_size <= 0:
raise ValueError(f"batch_size must be > 0: {self.batch_size}")
if self.epsilon < self.epsilon_min:
raise ValueError("epsilon must be >= epsilon_min")
if self.epsilon_decay < 0:
raise ValueError(f"epsilon_decay must be >= 0: {self.epsilon_decay}")
if self.gamma < 0:
raise ValueError(f"gamma must be >= 0: {self.gamma}")
@classmethod
def from_dict(cls, d: Mapping[str, Any]) -> "DQNConfig":
# 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) -> "DQNConfig":
# ignore unknown override keys as well
known_overrides = {k: v for k, v in overrides.items()
if k in self.__dataclass_fields__}
cfg = replace(self, **known_overrides)
cfg.validate()
return cfg
[docs]
class DQN_RC:
def __init__(self, env=None, *, reward_function=None, env_name=None, watch=True,
config: Optional[DQNConfig | Mapping[str, Any]] = None,
model: Optional[str] = None,
**overrides: Any) -> None:
"""
A class that implements a DQN Reservoir Computer for reinforcement learning tasks.
Args:
env (gym.Env): The environment to train the model on.
config (DQNConfig | Mapping[str, Any]): Hyperparameter values as a DQNConfig object or a dict. If both `config` and `overrides` are provided, `overrides` takes precedence.
model (str): Optional path to a previously trained model to load.
env_name (str): Optional name of the environment to create if `env` is not provided. Ignored if `env` is provided. Required if `env` is not provided.
watch (bool): If True, render the environment during training. Ignored if `env` is provided, as rendering should be set when creating the environment.
overrides (dict): Individual hyperparameter values
Attributes:
env (gym.Env): The environment to train the model on.
config (DQNConfig): Hyperparameter values including defaults, values from `config`, and any overrides.
rc_seed (int): Seed for random number generator for reservoir computer.
mask_seed (int): Seed for random number generator for input mask.
env_seed (int): Seed for random number generator for environment.
weight_seed (int): Seed for random number generator for readout weights.
rand (np.random.RandomState): Random number generator for reservoir computer.
N (int): Number of neurons in the reservoir.
theta (float): neuron separation time (relative to natural frequency of reservoir).
memory (deque): Memory buffer for experience replay.
learning_rate (float): Learning rate for training the readout weights.
learning_rate_decay (float): Decay factor for learning rate after each episode.
learning_rate_min (float): Minimum learning rate.
epsilon (float): Initial epsilon for epsilon-greedy policy.
epsilon_min (float): Minimum epsilon for epsilon-greedy policy.
epsilon_decay (float): Decay factor for epsilon after each episode.
gamma (float): Discount factor for future rewards.
trials (int): Number of training trials.
val_size (int): Number of episodes to run for validation.
target_update_rate (int): Number of trials between updates to the target readout weights.
state_shape (tuple): Shape of the environment's observation space.
action_shape (int): Number of actions in the environment's action space.
reward_function (callable): Optional function for shaping rewards.
loss (float): The most recent loss value from training the readout weights.
h (float): Integration step size for simulating the reservoir dynamics.
tau (int): Time delay for feedback in the reservoir (in number of periods, theta).
fb_gain (float): Feedback gain for the reservoir.
"""
if env is None:
mode = None
if watch:
mode = "human"
if env_name is not None:
env = gym.make(env_name, render_mode=mode)
else:
print("Warning: No environment provided. Defaulting to CartPole-v1.")
env = gym.make("CartPole-v1", 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.")
base = DQNConfig(env_name=self.env_name)
# merge user config (dataclass or dict)
if isinstance(config, DQNConfig):
base = config
elif isinstance(config, Mapping):
base = DQNConfig.from_dict(config)
elif config is not None:
raise TypeError("config must be a DQNConfig, dict-like, or None")
# apply overrides (highest priority)
self.config = base.updated(**overrides)
# future generalization addition
# if self.config.NormalizationFactor is None and env.__class__.__name__ == '':
# self.config.NormalizationFactor
# ----------------------------------------------------------------------------------------
### 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.rc_seed)
self.N = self.config.N # N^2 = number of internal nodes in the reservoir
self.env = env # ale environment
self.test_env = None
self.bufferLength = self.config.bufferLength # memory length, not sure why N is used -> in progress learning
self.memory = deque(maxlen=self.bufferLength) # memory buffer in form of a queue for experience replay
self.trials = self.config.trials
self.val_size = self.config.val_size
self.learning_rate = self.config.learning_rate # learning rate of training
self.learning_rate_decay = self.config.learning_rate_decay # learning rate decay after each episode, 1 means no decay
self.learning_rate_min = self.config.learning_rate_min # minimum learning rate
self.epsilon = self.config.epsilon # initial epsilon for epsilon-greedy policy
self.epsilon_min = self.config.epsilon_min
self.epsilon_decay = self.config.epsilon_decay
self.gamma = self.config.gamma # discount factor for future rewards
self.rewardNormalizationFactor = self.config.rewardNormalizationFactor # normalization factor for reward if needed
self.target_update_rate = self.config.target_update_rate # when to update the target method
self.batch_size = self.config.batch_size # batch size for training
self.update_counter = 0 # keep track of episodes for target model updating
self.state_shape = self.env.observation_space.shape
if self.state_shape is None:
print(self.env.observation_space)
self.state_shape = [len(self.env.observation_space)]
print(self.state_shape)
try:
self.action_shape = self.env.action_space.n
except:
self.action_shape = self.env.action_space.shape[0]
print(f'action shape: {self.action_shape}')
self.loss = 0 # Initialize loss to 0
if reward_function is None:
def r_func(*, reward, cur_state=None, new_state=None, action=None, done=None, hparams=None, validate=False):
return reward
self.reward_function = r_func
else:
self.reward_function = reward_function
# ----------------------------------------------------------------------------------------
### MEMS RC parameters
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)
if self.config.tau_N != 0:
self.tau = int(self.config.tau_N * self.N) # craft tau based on N if specified
print('setting tau based on N: ', self.tau)
self.NormalizationFactor = self.config.NormalizationFactor
self.NormalizationOffset = self.config.NormalizationOffset
self.amplification = self.config.amplification
# ----------------------------------------------------------------------------------------
# Code Optimization Metrics
self.mems_sim_times = np.array([])
# ----------------------------------------------------------------------------------------
# check if model is provided after defining all hyperparams
if model is not None:
print('transferring over model')
root, extension = os.path.splitext(model)
if extension.lower() != '.npz':
raise TypeError('File extension must be .npz')
self._load_reservoir(model)
self.W_out_target = self.W_out.copy()
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.mask_seed,
normalize_mask=True
)
weight_rand = np.random.RandomState(self.weight_seed)
scale = 1/np.sqrt(self.N)
self.W_out = weight_rand.normal(0, scale, size=(self.N+1+self.state_shape[0], self.action_shape))
self.W_out_target = self.W_out.copy()
def _readout(self, neurons, W_out, *, analyze=False):
"""
Calculates the best action based on neuron states and weights.
Args:
neurons (np.ndarray): The current state of the neurons.
W_out (np.ndarray): Weight matrix for the output layer.
analyze (bool): If True, returns the full list of Q-values calculated.
Returns:
high_Q (float): The maximum Q-value found.
action (int): The index of the winning action.
Q_vals (list[float] | None): A list of all Q-values if `analyze` is True, otherwise None.
"""
high_Q = float('-inf')
Q_vals = np.array([])
for i in range(W_out.shape[1]):
Q_val = np.dot(neurons, W_out[:,i])
if analyze:
Q_vals = np.append(Q_vals, Q_val)
if Q_val > high_Q:
high_Q = Q_val
action = i
if analyze:
return high_Q, action, Q_vals.tolist()
else:
return high_Q, action, None
def _optimize(self, target, Q_current, neurons, W_old): #There will be a different readout circuit training for each action
"""
Run one step of gradient descent to update the readout weights.
Args:
target (float): Target Q-value.
Q_current (float): Current Q-value prediction.
neurons (np.ndarray): Reservoir neurons state including input feedback.
W_old (np.ndarray): Current readout weights to be updated.
Returns:
W_updated (np.ndarray): Updated readout weights after one step of gradient descent.
loss (float): The squared difference between target and predicted Q-value.
"""
loss = np.square(target - Q_current)
# Gradient = (-2) * np.multiply(np.transpose(Z), (target - Q_current)).squeeze()
Gradient = -2 * (target - Q_current) * neurons.flatten()
W_updated = W_old - self.learning_rate * Gradient
# print(f'max_w_old: {np.max(W_old)}, max_w_new: {np.max(W_updated)}, Gradient: {np.max(Gradient)}, learn_rate: {self.learning_rate}')
return W_updated, loss
def _act(self, state, opt=False, *, analyze=False, mode='train'):
"""
Run the reservoir computer to find next action and apply epsilon-greedy policy.
Args:
state (np.ndarray): Current state of the environment.
opt (bool): If True, applies epsilon-greedy policy and decays epsilon.
analyze (bool): If True, returns the full list of Q-values calculated.
Returns:
action (int): Chosen action.
neurons (np.ndarray): Reservoir neurons state including input feedback.
Q_vals (list[float] | None): A list of all Q-values if `analyze` is True, otherwise None.
"""
if opt:
self.epsilon *= self.epsilon_decay
self.epsilon = max(self.epsilon_min, self.epsilon)
# Action_out = self.predictRC(state)
time_start = time.time_ns()
neuron_vals = self.reservoir.sim(obs=state, direct_fb=True, mode=mode) # shape (N,)
self.mems_sim_times = np.append(self.mems_sim_times, time.time_ns() - time_start) if opt else self.mems_sim_times
self.MEMS_neurons = np.array(neuron_vals).reshape(1, -1) # shape (N, 1)
_, action, Q_vals = self._readout(self.MEMS_neurons, self.W_out, analyze=analyze)
if opt:
if self.rand.random() < self.epsilon:
action = int(self.rand.randint(0,self.action_shape)) # using rc_seed instead of env
return action, self.MEMS_neurons, Q_vals
# return self.env.action_space.sample(), self.MEMS_neurons
return action, self.MEMS_neurons, Q_vals
def _remember(self, action, reward, done, neurons, future_neurons):
"""
Add experience (action, reward, done, neurons, future_neurons) to the memory buffer
for experience replay-based training.
"""
self.memory.append([action, reward, done, neurons, future_neurons])
def _replay(self):
"""
Optimizes readouts based on a batch of experiences from memory.
1. Takes a random batch of experiences from memory.
2. For each experience, runs target reservoir computer to get target Q-value.
3. Trains the readout matrix (optimize) depending on what action the actual model took.
"""
self.loss = 0
batch_loss = 0 # Initialize batch loss
if (len(self.memory) < self.batch_size):
return
# changed the sampling technique to allow for the same random number generator to be used
sample_indices = self.rand.choice(len(self.memory), self.batch_size, replace=False)
samples = [self.memory[i] for i in sample_indices]
for sample in samples:
action, reward, done, neurons, future_neurons = sample
# Do things related to the readout circuit of action 0
if done:
target = reward/self.rewardNormalizationFactor
# print(f'target: {target}, normfactor: {self.rewardNormalizationFactor}')
else:
Q_future = self._readout(future_neurons, self.W_out_target)[0]
target = (reward + Q_future * self.gamma)/self.rewardNormalizationFactor
# print(f'target: {target}, normfactor: {self.rewardNormalizationFactor}')
Q_current = np.dot(neurons, self.W_out[:,action])
new_weights, loss = self._optimize(target, Q_current, neurons, self.W_out[:,action])
self.W_out[:,action] = new_weights
batch_loss += loss # Accumulate loss for the batch
# Decay learning rate
self.learning_rate *= self.learning_rate_decay
self.learning_rate = max(self.learning_rate, self.learning_rate_min) #Annealing
# Update target network every few episodes
self.update_counter = self.update_counter+1
if(self.update_counter%self.target_update_rate == 0):
self.W_out_target = self.W_out.copy()
average_batch_loss = batch_loss / self.batch_size # Calculate average batch loss
self.loss = average_batch_loss # Return average batch loss
def _play_env(self, obs, env, *, opt=False, mode='train'):
done = False
total_reward = 0
trial_length = 0
trial_loss = 0
tot_neuron_sat = 0
self.reservoir._zero_reservoir(mode=mode)
action, neurons, _ = self._act(obs, opt, mode=mode)
while not done:
new_state, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated
# ------------------------------------------------------------------
# Reward Shaping (if necessary)
# reward = reward + 1
if opt:
reward = self.reward_function(reward=reward, cur_state=obs, new_state=new_state, action=action, done=done, hparams=self.config)
obs = new_state
# ------------------------------------------------------------------
future_action, future_neurons, _ = self._act(new_state, opt)
if opt:
self._remember(action, reward, done, neurons, future_neurons)
self._replay()
trial_loss += self.loss
neurons = future_neurons
action = future_action
total_reward += reward
tot_neuron_sat += np.sum(np.abs(neurons)>=0.89)*100/neurons.size
trial_length += 1
returns = {
"total_reward": total_reward,
"trial_length": trial_length,
"trial_loss": trial_loss if opt else None,
"avg_neuron_sat": tot_neuron_sat/trial_length
}
return returns
[docs]
def train(self, *, env=None, meta={}, trials=None,
folder_path=None, wandb_on=True, save_model=True,
ray_on=False, ray_metric='avg_reward', ray_mode='max',
ray_report_freq=1, full_validate=False, val_size=5,
val_size_min=1, threed_it=False, resume=False, wandb_id=None):
"""
Runs a reinforcement learning training on the given environment
for the initialized DQN_RC model.
Args:
env (gym.Env): The environment to train on (likely already defined during agent creation).
meta (dict): A dictionary containing metadata for the training session.
trials (int): The number of trials to train on.
folder_path (str): The path to save the model to.
wandb_on (bool): Whether to use wandb for logging.
save_model (bool): Whether to save the model to disk.
full_validate (bool): Whether to fully validate every trial or when an agent gets close to current best reward.
val_size (int): The number of environment seeds in the validation set.
val_size_min (int): The minimum number of samples in the validation set, if full_validate is False.
threed_it (bool): Used to create policy map (to be implemented more easily)
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 env is not None:
self.env = env
if trials is None:
trials = self.trials
required_fields = ["project", "group", "job_type", "run_name", "tags"]
for field in required_fields:
if field not in meta:
meta[field] = f"{field}_dne"
self.run_name = meta["run_name"]
if folder_path is not None:
self.folder_path = folder_path
elif save_model:
print("No folder provided. Model will be saved to 'DQN_RC/models' in current directory.")
self.folder_path = f"DQN_RC/models/{self.env_name}-{meta['group']}"
self.datetime_str = datetime.now().strftime('%Y%m%d-%H%M%S')
self.run_animal = self._get_animal()
if not resume:
meta["run_name"] = f'{self.run_name}-{self.run_animal}-{self.datetime_str}'
# setup wandb
if wandb_on:
run = wandb.init(
project=meta["project"],
group=meta["group"],
job_type=meta["job_type"],
name=meta["run_name"],
tags=meta['tags'],
config={"hp": self.config, "meta": meta},
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'
print(f'Training starting for run: {self.run_name}')
eval_metric = float('-inf')
if ray_on and ray_mode == 'min':
eval_metric = float('inf')
trial_times = np.array([])
if threed_it:
images =[]
best_avg_reward = float('-inf')
obs = self.env.reset(seed=self.env_seed)[0]
for trial in range(trials):
if trial != 0:
obs = self.env.reset()[0]
self.mems_sim_times = np.array([])
time_start = time.time_ns()
returns = self._play_env(obs, self.env, opt=True)
trial_loss = returns["trial_loss"]
total_reward = returns["total_reward"]
trial_length = returns["trial_length"]
val_time_start = time.time_ns()
if full_validate or ray_on:
val_results = self._validate(val_size=val_size)
if val_results['avg_reward'] >= best_avg_reward or threed_it:
if not ray_on:
print(f'New best reward: {val_results["avg_reward"]}')
best_avg_reward = val_results['avg_reward']
if save_model and not ray_on:
if resume:
name_val = f"resume_best_{self.datetime_str}"
else:
name_val = f"best_{self.run_animal}_{self.datetime_str}"
best_model_path = self._save_reservoir(name_val=name_val)
elif total_reward > .9*best_avg_reward or total_reward > 1.1*best_avg_reward: # validate fully if reward is close to best
print('sampling')
val_results = self._validate(val_size=val_size)
if val_results['avg_reward'] >= best_avg_reward :
print(f'New best reward: {val_results["avg_reward"]}')
best_avg_reward = val_results['avg_reward']
if save_model and not ray_on:
if resume:
name_val = f"resume_best_{self.datetime_str}"
else:
name_val = f"best_{self.run_animal}_{self.datetime_str}"
best_model_path = self._save_reservoir(name_val=name_val)
else: # if full_validation is false and not close to best, validate for val_size_min
val_results = self._validate(val_size=val_size_min)
if resume and not ray_on:
name_val = f"resume_last_{self.datetime_str}"
else:
name_val = f"last_{self.run_animal}_{self.datetime_str}"
if save_model and not ray_on:
self._save_reservoir(name_val=name_val)
val_time = time.time_ns() - val_time_start
if threed_it:
from analyze.analyze_run import AnalyzeRun
ar = AnalyzeRun(best_model_path)
img = ar.visualize_state_action(num_steps=10,
ele_offset=-10,
azim_offset=-90,
folder_path=f'{folder_path}/3d_trials',
file_name=f'3d_trial{trial}',
show=False,
# save=False,
# to_bytes=True,
title_addon=f'- Trial {trial}: reward={int(best_avg_reward)}')
# images.append(img)
# np.savez_compressed(f'{folder_path}/{meta["run_name"]}_images.npz',
# images=images,)
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(path_npz=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'],
'eval_metric': eval_metric,
'avg_trial_loss': trial_loss,
'eval_metric_cur': val_results[ray_metric],
'trial': trial,
'epsilon': self.epsilon,
'learning_rate': self.learning_rate,
'trial_length': trial_length,
}, 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'],
'eval_metric': eval_metric,
'avg_trial_loss': trial_loss,
'eval_metric_cur': val_results[ray_metric],
'trial': trial,
'epsilon': self.epsilon,
'learning_rate': self.learning_rate,
'trial_length': trial_length,
})
trial_time = time.time_ns() - time_start
trial_times = np.append(trial_times, trial_time/1_000_000_000) # should be in seconds
if wandb_on:
wandb.log({
"reward": total_reward,
"avg_loss_this_trial": trial_loss/trial_length,
"trial_length": trial_length,
"epsilon": self.epsilon,
"learning_rate": self.learning_rate,
"Neuron Saturation Percentage": val_results['avg_neuron_sat'],
"Average Trial Time per N [ms]": trial_time/trial_length/1000000/self.N,
'val_reward': val_results['avg_reward'],
'best_avg_reward': best_avg_reward,
'trial': trial,
})
if not ray_on:
print(f"Trial {trial+1}/{trials}, Trial Reward: {total_reward:.1f}, best_avg_reward: {best_avg_reward:.1f}, Trial Length: {trial_length}, Avg Loss: {np.max(trial_loss/trial_length):.2f}, Avg neuron sat: {val_results['avg_neuron_sat']:.1f} [%]")
print(f"Trial time = {trial_time/1000000000:.2f} [s], MEMS sim time per node = {np.mean(self.mems_sim_times)/self.N/1000000:.3f} [ms], Validation time = {val_time/1000000000:.2f} [s]")
hours = np.floor(np.sum(trial_times)/60/60)
if not ray_on:
print(f'done, trained for {int(np.sum(trial_times))} [s] or {int(hours)} hours and {np.sum(trial_times/60)-hours*60:.2f} minutes')
if wandb_on and not ray_on:
wandb.finish()
if ray_on:
return
# elif threed_it:
# return image_path
else:
return best_model_path
def _validate(self, val_size=None):
# if self.test_env is None:
# print(self.env_name)
# if self.env_name == 'tron':
# from games.tron import TronEnv
# self.test_env = TronEnv()
# else:
self.test_env = gym.make(f'{self.env.unwrapped.spec.id}')
avg_neuron_sats = []
rewards = []
trial_lengths = []
trial_losses = []
self.reservoir._zero_reservoir(mode='val')
if val_size is None:
val_size = self.val_size
start = 20
end = start + val_size
for i in range(start,end,1):
obs = self.test_env.reset(seed=i)[0]
returns = self._play_env(obs, self.test_env, mode='val')
rewards.append(returns["total_reward"])
trial_lengths.append(returns["trial_length"])
avg_neuron_sats.append(returns["avg_neuron_sat"])
val_results = {
'avg_reward': np.mean(rewards),
'avg_trial_length': np.mean(trial_lengths),
'avg_neuron_sat': np.mean(avg_neuron_sats),
}
return val_results
def _save_reservoir(self, *, name_val=None, file_path=None):
"""
Save reservoir to disk as a compressed NPZ file, including metadata about the model configuration and training.
"""
self.hparams = {
'algorithm': 'DQN_RC',
'environment': self.env.unwrapped.spec.id if self.env_name != 'tron' else 'tron',
'run_name': self.run_name
}
self.hparams.update(asdict(self.config)) # Convert dataclass fields to dict and update hparams
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,
reservoir=self.reservoir,
W_out=self.W_out,
meta=meta)
def _load_reservoir(self, path_npz):
with np.load(path_npz, allow_pickle=True) as data:
mask = data['mask'] if 'mask' in data else None
self.reservoir = data['reservoir'].item() if 'reservoir' in data else None
self.W_out = data['W_out']
# 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 not hasattr(self, 'reservoir') or 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)
# Store as a Python dict
self.hparams = meta_dict
def _get_animal(self):
"""
Returns the name of a random animal for run naming.
"""
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']