myriad-jax 0.1.0

Platform for large-scale high-throughput control & decision-making experiments

Myriad

JAX-native platform for massively parallel control, system identification, and active learning of uncertain, stochastic systems.

[!WARNING] Myriad is in early active development — we're building in public. APIs will change, documentation has gaps, and some features are still taking shape. Things will improve over time. Contributions, feedback, and ideas are very welcome — open a discussion or reach out to Robin ([e-mail]).

Jump straight to Installation or Quickstart to see Myriad in action, or check out the [full documentation](todo: add link). 🤸🏾

TODO: add 2 videos here: (a) learned agent on ccas-ccar, (b) real-time system ID via active learning.

At a Glance

Myriad is a playground to explore RL, traditional control, system identification, and active learning — with a focus on problems where uncertainty, stochasticity, and rare discrete dynamics play a big role and force us to study very large numbers of variants in parallel (think: biology → system = cell, chemistry → system = reactor, THIRD EXAMPLE?). 🛝

It's a ready-to-go experimental platform. You can use one of the already-implemented tasks/problems, algorithms, or implement your own and simply plug them in. Myriad will handle the intricacies of JAX/GPU optimization, training/evaluation loops, hyperparameter tracking, metrics logging, and many more not-so-fun things — freeing time for the more fun science and engineering bits. 👩🏾‍🔬👨🏻‍🔬

Last but not least, it yields results that are 100% reproducible. 🌟

Interested in the story behind Myriad? Read our [Mission Statement & Philosophy](TODO: LINK_TO_DOCS_HERE).

Key Features

• ⚡ Massive GPU Parallelism: run algorithms on 1M+ of environments simultaneously.

• 🏎️ JAX JIT Optimization: Myriad is fast, even on CPU.

• ✅ 100% Reproducible: Myriad is fully deterministic. Using the same initial random seed and configuration file will yield the same results → great for science.

• 🎲 Exact Stochastic Simulations: native JAX implementation of the Gillespie Algorithm (aka SSA) for discrete, asynchronous molecular events.

• ∇ Differentiable "White-Box" Physics: exposes underlying physics, ODEs, and jump processes for gradient-based system ID and active learning.

• 🛠 Research-Ready: pre-configured with Hydra, Pydantic, and W&B support.

Ecosystem

Many amazing RL x JAX tools already exist! Here's how we believe Myriad complements them.

Feature	Gymnasium/Gymnax	Brax	Myriad
Best For	Standard RL benchmarks	Robotics & Locomotion	Wet-Lab / Scientific Systems
Physics	Black Box / Various	Rigid Body (Contacts)	Stochastic, ODEs, Jump Processes
Differentiable?	No	Yes	Yes
System ID	Low support	Low support	Key focus
Primary Goal	Agent Performance	Fast Physical Control	Active Learning & Stochastic Control

• Use Myriad if: you model biological/chemical systems, require Gillespie/SSA stochasticity, or need active learning for parameter uncertainty.

• Use Brax if: you need massive-scale robotics or contact dynamics.

• Use Gymnasium, Gymnax or JaxMARL if: you need standard baselines or multi-agent RL.

Installation

Requirements: Python 3.10+, JAX 0.7+.

[!IMPORTANT] GPU Support: JAX installation can be hardware-specific. We strongly recommend installing JAX according to your CUDA/cuDNN version before installing Myriad if you encounter issues.

With pip

# Standard installation
pip install myriad-jax

# With generic GPU dependencies (checks for nvidia-related packages)
pip install "myriad-jax[gpu]"

From source (for development)

git clone https://github.com/robinhenry/myriad-jax.git
cd myriad-jax
poetry install --with dev,gpu

Quickstart

Myriad is designed to be used programmatically (for research loops) or via CLI (for massive sweeps).

Python API

from myriad import create_config, train_and_evaluate

# Configure a gene expression control experiment across 10k cells
config = create_config(
    env="gene-circuit-v1",      # A stochastic gene circuit (Gillespie)
    agent="dqn",                # The algorithm to use (eg, 'pid', 'pqn')
    num_envs=10_000,            # 10k parallel simulations (cells)
    scan_autotune=True,         # Automatically optimize GPU training loop parameters
)

# Run the experiment (JIT-compiled & distributed on GPU)
results = train_and_evaluate(config)

# Inspect performance metrics
metrics = results.eval_metrics
print(f"Return (mean +/- std): {metrics.mean_return} +/- {metrics.std_return}")

# Quickly plot convergence curve
TO ADD

CLI Usage

Leverage Hydra to run massive parameter sweeps or experiments directly from the terminal.

# Train a DQN agent on 50,000 parallel cartpole environments
myriad train env=cartpole-control run.num_envs=50000 agent=dqn

# Switch to System Identification mode (inferring hidden physics parameters)
myriad train env=gene-circuit-v1 agent=pqn sysid=true

See the [Documentation](add link) for a full list of examples and configuration overrides.

Flagship Environments

To add.

---

A little bit of history: Myriad is named after the Greek myrias ("ten thousand"), inspired by microfluidic "mother machines" that observe 100,000+ cells simultaneously. It brings this paradigm to computational research: providing a myriad of viewpoints from which to learn about and control complex systems — whether they are biological circuits, chemical reactors, or robotic swarms.