hardware_rc.ppo_rc

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.

Classes

PPOConfig(env_name, general_seed, mask_seed, ...)

PPO_RC([env, reward_function, env_name, ...])

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)[source]

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)[source]

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)