hardware_rc package

hardware_rc.analyze module

hardware_rc.dqn_rc module

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.

class hardware_rc.dqn_rc.DQN_RC(env=None, *, reward_function=None, env_name=None, watch=True, config: DQNConfig | Mapping[str, Any] | None = None, model: str | None = None, **overrides: Any)

Bases: object

train(*, 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.

Parameters:

• 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:

The path to the best model found during training.

Return type:

best_model_path (str)

hardware_rc.ppo_rc module

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.

class hardware_rc.ppo_rc.PPO_RC(env=None, *, reward_function=None, env_name=None, watch=True, N_envs=10, config: PPOConfig | Mapping[str, Any] | None = None, model: str | None = None, **overrides: Any)

Bases: object

train(*, 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.

Parameters:

• 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:

The path to the best model found during training.

Return type:

best_model_path (str)

hardware_rc.reservoir module

A module to simulate the dynamics of a reservoir computer with an RK4 DDE solver along with a analysis class for visualizing and testing reservoirs.

It is built with JAX, specifically JIT compilation, for speed.

class hardware_rc.reservoir.AnalyzeReservoir(reservoir)

Bases: object

bode_plot(*, N=80, amplification=5, VAC=1, delay=35, folder_path, file_name, show=True, save=False, fig_size=(5, 3.5))

create_sine_input(steps, VAC, N)

find_ss_amp(positions, delay)

plot_Vin(*, title_addition='', save=False, folder_path=None, base_file_name='', figsize=(3.5, 3), ax=None, show=True, color='black', state=[0, 0, 0, 0], Vin=None, partial_N=None, state2=None, color2='orange')

plot_mask(*, together=True, separate=True, combined=True, cmap_color='tab10', title_addition='', save=False, folder_path=None, base_file_name='', figsize=(3.5, 3), ax=None, show=True, color='royalblue', partial_N=None)

plot_response(data, *, points=None, title=None, save=False, folder_path=None, base_file_name='', overlay_mask=False, Vin=None, show=True, figsize=(10, 5), show_neuron_sat=False, neuron_sat=None, reps, return_image=False, ax=None, color='royalblue', color2=None, x_on=True, one_y=False, conf=False, Vin_color='green', show_points=True, adjust_dict=None, xticks=None, xtick_labels=None)

reservoir_response_subplots_2_mc(*, figsize=(3.5, 2.5), dpi=150, states=None, show=True, save=True, folder_path=None, file_name=None)

reservoir_response_subplots_3(*, figsize=(3.5, 2.5), dpi=150, states=None, show=True, save=True, folder_path=None, file_name=None)

reservoir_response_subplots_3_mc(*, figsize=(3.5, 2.5), dpi=150, states=None, show=True, save=True, folder_path=None, file_name=None)

reservoir_response_subplots_4(*, figsize=(3.5, 2.5), dpi=150, states=None, show=True, save=True, folder_path=None, file_name=None)

reservoir_response_subplots_4_mc(*, figsize=(3.5, 2.5), dpi=150, states=None, show=True, save=True, folder_path=None, file_name=None)

reservoir_subplots(*, figsize=(3.5, 2.5), dpi=150, states=None, show=True, save=True, folder_path=None, file_name=None)

sim_response(*, state=None, seed=None, overlay_mask=False, overlay_Vin=False, overlay_actual_Vin=True, title=None, save=False, folder_path=None, base_file_name='', show=True, figsize=(10, 5), show_neuron_sat=False, reps=1, return_image=False, states=None, ax=None, color='royalblue', color2='royalblue', x_on=True, one_y=False, Vin_color='green', show_points=True, adjust_dict=None, xticks=None, xtick_labels=None)

class hardware_rc.reservoir.Reservoir(*, h=0.02, theta, N, sd=0, amp=1, tau=0, fb_gain=0, norm_factor=None, norm_offset=None, state_shape=1, input_connectivity=0.2, mask=None, mask_seed=0, load=False, normalize_mask=True, VDC=20, n_envs=1, n_steps_needed=None)

Bases: object

create_mask(seed, normalize_mask=False)

static diff_eq(y, V, amp, VDC)

MEMS differential equation.

Parameters:

• y (jnp.ndarray) – Current state of the MEMS device [position, velocity].

• V (jnp.ndarray) – Input voltage at current time [Volts].

• amp (jnp.float) – Amplification parameter for the reservoir.

static rk4(*, n_steps, pos_buf, vel_buf, h, Vin, theta, N, sd, amp, stride, tau, fb_gain, buf_idx, buf_cnt, VDC)

Simulate and sample a DDE for a given set of step input voltages with RK4 and cubic Hermite interpolation for the delay term. The state of the DDE is stored in ring buffers to allow for efficient and JAX correct sampling and delayed feedback.

Parameters:

• n_steps (int) – total number of simulation steps.

• pos_buf (jnp.ndarray) – ring buffer to store past positions, length buf_len.

• vel_buf (jnp.ndarray) – ring buffer to store past velocities, length buf_len.

• h (float) – step size for the simulation.

• Vin (jnp.ndarray) – input voltage array, length n_steps.

• theta (float) – sample time for one reservoir reading [non-dimensional].

• N (int) – number of reservoir states.

• sd (int) – sample delay (in timesteps).

• stride (int) – sample every ‘stride’ with delay ‘start’.

• tau (float) – delayed feedback gain parameter [non-dimensional].

• fb_gain (float) – feedback gain parameter [non-dimensional].

• amp (float) – amplitude of the input voltage [non-dimensional].

• buf_idx (int) – next write index (most recent sample is at buf_idx-1).

• buf_cnt (int) – number of valid samples in buffers (<= len(pos_buf)).

• VDC (float) – DC offset voltage [V].

Returns:

all positions from simulation. sampled (jnp.ndarray): sampled positions (neuron values), length n_samples. pbuf_out (jnp.ndarray): updated ring buffer of past positions. vbuf_out (jnp.ndarray): updated ring buffer of past velocities. buf_idx_out (int): updated next write index. buf_count_out (int): updated number of valid samples in buffers.

Return type:

positions (jnp.ndarray)

sim(*, obs=None, Vin=None, mode='train', env_idx=0, full_data=False, direct_fb=False, get_Vin=False, update_state=True)