API reference

Every public module, class and function in the coadapt package, with its docstring pulled straight from the source. Read them in pipeline order: spaces → env → envs → network → agent → rollout → solvers → storage → plotting.

Modules

coadapt

coadapt — a small, decentralized multi-agent RL harness under partial observability.

Pipeline (read the modules in this order to follow the flow):

    spaces   -> what an action looks like (Discrete / MultiDiscrete / Box)
    env      -> the Env contract + EnvSpec (the base class)
    envs/    -> the N-actor worlds: task state, per-actor measurement, team reward
    network  -> the communication layer: who hears whom, and how stale (partial obs)
    agent    -> N independent local policies (one per actor), no parameter sharing
    rollout  -> the observe -> act -> step -> communicate loop (scan over time)
    solvers  -> from-scratch REINFORCE (one independent learner per agent)
    storage  -> save a run as a structured dict file
    plotting -> read those dict files and draw the standard figures

The whole point of the design: environment actor i is bound 1:1 to agent i, each agent
sees only its own measurement plus *delayed* messages from graph neighbours, and each
agent trains its own policy — a decentralized, partially-observed cooperative setting.

coadapt.agent

agent.py — N independent local policies, one per actor, with NO parameter sharing.

Each agent is a small memoryless MLP:  obs -> hidden (tanh) -> action-distribution params.
The obs an agent sees is its OWN measurement (from the env) concatenated with the DELAYED
message from its parent (from the network). That concatenation is the only thing the policy
ever sees — partial observability is therefore structural.

How we keep N *independent* policies while staying fast on a GPU:
  - We store all N policies as ONE pytree whose leaves have a leading agent axis [N, ...].
    Agent i's parameters are just the i-th slice of every leaf.
  - `jax.vmap(single_policy_forward)(params, obs)` runs all N MLPs at once, each on its own
    parameters and its own obs. Same math as a Python loop over agents, but vectorised.
  - Because every agent has its own leaf slice, updating agent i never touches agent j
    (guaranteed; there is no shared weight tensor). This is what "strictly decentralized"
    means here, in contrast to the old shared-parameter IPPO.

The `Agent(params, i)` handle is a thin convenience view for the experimenter / tests: it
exposes agent i's own parameters. The hot path uses the vmapped functions directly.

class Agent(self, agents, i)

A read-only view of ONE agent inside the stacked bundle. For tests / introspection.

`agents` is the stacked pytree; `i` is the agent index. `.policy` returns agent i's own
parameter pytree (its leaf slices), demonstrating that each agent's policy is local.
forward(self, obs_i)
Run just this agent's policy on a single obs [obs_dim] -> dist params.

act(agents, obs, key, action_space)

All N agents act at once on their own obs (the decentralized-execution step).

agents : stacked pytree, leaves [N, ...]
obs    : [N, obs_dim]  each agent's own (partial) observation
returns:
    actions      : env-native action array (shape depends on the space)
    dist_params  : [N, param_size]  each agent's raw distribution params (also the message)
    logp         : [N]  log-prob of the action each agent took

forward(params, obs)

One agent: obs [obs_dim] -> action-distribution params [param_size].

init_agents(key, N, obs_dim, param_size, hidden=32)

Build N INDEPENDENT policies as one stacked pytree with leading axis [N, ...].

We split the key N ways and vmap init_params, so each agent gets its own random weights.

init_params(key, obs_dim, param_size, hidden=32)

Random parameters for ONE agent's MLP: obs_dim -> hidden -> param_size.
Plain dict of arrays (a pytree). Small Glorot-ish init.

coadapt.env

The environment contract: EnvSpec + the base class every world implements.

An environment here is ONLY the task/world dynamics:
  - it holds the hidden target the team must match,
  - it produces each actor's own *internal measurement* (what agent i physically senses),
  - it turns the agents' actions into a shared team reward.

It does NOT handle communication — that is the job of network.py. Keeping the two apart is
what makes "partial observability lives in the network" true and testable.

The world is functional (JAX-friendly): reset/step take a key and a state and return new
values instead of mutating anything, so we can `scan` over time and `vmap` over seeds.

    reset(key)                -> (measurement, world)
    step(key, world, actions) -> (measurement, world, reward, done, info)

    measurement : float32 [N, meas_dim]   what each actor senses of ITSELF/the world
    world       : a small frozen dataclass holding the hidden task state
    reward      : scalar shared team reward in [0, 1]
    actions     : shape depends on the action space (see spaces.py)
    info        : dict; MUST contain "per_actor_error" -> float32 [N] (per-actor tracking
                  error). The rollout records it for the error-vs-distance metric, so every
                  env's step must return this key. Additional keys are free-form.

Reference bound note: `solvers.random_oracle.oracle_return` assumes the relay measurement
layout (measurement[:,0]=is_source flag, measurement[:,1]=target-if-source). A new world with
a different measurement layout should provide its own oracle/ceiling rather than rely on that
one (it raises TypeError for action spaces it does not cover).

class Env

Base world. Subclasses set `self.spec` and implement reset/step.
reset(self, key)

(no docstring)

step(self, key, world, actions)

(no docstring)

class EnvSpec(self, num_agents: int, meas_dim: int, action_space: Any, horizon: int, jittable: bool = True) -> None

Everything a solver needs to know about a world without looking inside it.

coadapt.envs

Concrete worlds, each implementing the Env contract in ../env.py.

Stage 1 ships two:
  relay_discrete  — discrete actions (the Phase-1 reference-tracking toy)
  robot_consensus — continuous actions (a leader-follower platoon)

Later stages add the checklist (multi-discrete) and coordination worlds.

coadapt.envs.relay_discrete

relay_discrete — the discrete reference-tracking world (the Phase-1 toy, in JAX).

The story in one paragraph:
  There is a hidden scalar target theta_t in [0,1] that the whole team must match. Only the
  SOURCE actor (actor 0) senses theta_t directly. Every other actor senses only its own last
  action. So a non-source actor can only learn theta_t from what its neighbours tell it
  (handled by network.py) — and because messages are delayed, information about a drifting
  target arrives late the further you are from the source. The team reward rewards everyone
  for matching theta_t. When theta_t drifts (nu > 0), distant actors chase a stale value and
  the reward collapses. When theta_t is fixed (nu = 0), it is an easy static task (the control).

This file owns ONLY the target, the drift, the per-actor measurement, and the reward. Who
hears whom (and how stale) is network.py's job.

Action space: Discrete(K)  — each actor outputs an integer in {0..K-1}; a/(K-1) is its guess
of theta_t in [0,1].

class RelayDiscreteEnv(self, N=8, K=5, horizon=60, nu=0.0, source=0)

(no docstring)

reset(self, key)

(no docstring)

step(self, key, world, actions)

(no docstring)

class RelayWorld(self, theta: jax.Array, last_action: jax.Array, t: jax.Array) -> None

The hidden task state. A chex.dataclass so it is a JAX pytree we can scan over.

coadapt.envs.robot_consensus

robot_consensus — the continuous world (a leader-follower platoon).

The story:
  N robots on a line. The leader (actor 0) knows a target position ref. Every other robot
  wants to sit at ref too, but only senses the robot ahead of it (through the network, so
  possibly delayed). Actions are continuous positions in [low, high]. The reward is how
  close the whole platoon is to ref. This is the continuous twin of relay_discrete: same
  "information arrives late down the line" story, but with a continuous action so we exercise
  the Box action space and a Gaussian policy.

Owns ONLY the target, its drift, the per-robot measurement, and the reward. Communication
(who senses whom, delayed) is network.py's job.

Action space: Box((1,), low, high) — each robot outputs its next position.

class RobotConsensusEnv(self, N=8, horizon=60, nu=0.0, source=0, low=-1.0, high=1.0)

(no docstring)

reset(self, key)

(no docstring)

step(self, key, world, actions)

(no docstring)

class RobotWorld(self, ref: jax.Array, pos: jax.Array, t: jax.Array) -> None

RobotWorld(ref: jax.Array, pos: jax.Array, t: jax.Array)

coadapt.network

network.py — the communication layer, and the ONLY place partial observability lives.

What this module does:
  - builds a communication graph over the N actors (line / ring / star / grid / tree),
  - from a chosen source, computes each actor's PARENT (the neighbour one hop toward the
    source) and DISTANCE (hops to the source),
  - keeps a delay buffer of recent MESSAGES so that what an actor hears from its neighbour
    is `d` rounds stale (messages arrive late; information from far away arrives later),
  - exposes the two pings the plan requires:
        ping_channel(net, i)  -> what agent i is ALLOWED to hear (delayed neighbour message)
        ping_truth(net,...)   -> ground truth, for the experimenter's logs/plots ONLY.

A "message" from actor j is the concatenation of:
    [ state-feature(j) , action-distribution-params(j) ]
i.e. what j did (its normalized action) and how j's policy is currently behaving. Agents
communicate BOTH, per the design. Messages are compact float vectors so the delay buffer
stays small.

The key modelling choice (kept simple on purpose): an actor's channel observation is the
delayed message of its PARENT. On a line this reproduces the classic result that information
about a change is stale by ~ distance x (1 + delay) — the "information-arrival wall". The
source has no parent, so it hears nothing over the channel (it already senses the target).

class NetState(self, parent: jax.Array, dist: jax.Array, buffer: jax.Array) -> None

Everything about communication. The static graph fields never change during a run;
only `buffer` advances each round. Kept in one pytree so it scans cleanly.

class Network(self, N, action_space, topology='line', source=0, delay=0)

Holds the graph shape + message size and knows how to build/advance NetState.

`msg_dim = state_feature_size + action_param_size`. We fix it at construction so the
delay buffer has a static shape (required by JAX).
init_state(self)
A fresh NetState with an empty (zero) message buffer.
make_message(self, action_feature, action_params)
Pack one round of messages: [N, msg_dim] = concat(feature, dist-params).
ping_channel(self, net)
What every agent is ALLOWED to hear this round: the DELAYED message of its parent.

Returns [N, msg_dim]. The source hears a zero message (it has no parent / no need).
This is the only communication an agent's observation may use — enforced by tests.
ping_truth(self, net, message)
GROUND TRUTH for the experimenter only: the *current* (undelayed) message of
every actor. Never wire this into an agent's observation — it would erase partial
observability. Used by storage/plotting to see what actually happened.
push(self, net, message)
Advance the delay buffer by inserting this round's message at the front and
dropping the oldest. buffer[0] is newest; buffer[delay] is `delay` rounds old.

coadapt.plotting

plotting.py — the plotting harness: read stored run dicts and draw standard figures.

This module is a PURE CONSUMER of storage.py output. It never runs a simulation; it only
turns saved numbers into pictures. That separation means a figure can always be regenerated
from a saved run without recomputing anything.

Each function takes a loaded run dict (see storage.load_run) and an output path, draws one
sensible figure, and returns the path written.

plot_error_by_distance(run, path)

Per-agent tracking error vs hops-from-source (the information-arrival wall).
Expects curves['error_by_distance'] = [max_dist+1] and curves['dist'] optional.

plot_learning_curve(run, path)

A learning curve: mean per-step team reward vs iteration, with random/oracle lines
if they were recorded in metrics. `run` is a loaded run dict.

plot_surface_3d(run, path)

3-D collapse surfaces: z = achievable (oracle) reward over (delay, nu), one subplot
per N. Shows how the surface sinks as the network grows — the information-arrival wall.

Expects curves['delays'], curves['nus'], and one Z array per N named 'oracle_N<k>' of
shape [len(nus), len(delays)]. config['Ns'] lists the N values (subplot order).

plot_sweep(run, path)

Consolidated sweep figure: mean per-step team reward vs N, one line per
(condition, solver). Expects curves['N'] and curves['<condition>_<solver>'] arrays.
`run['config']['series']` lists the (condition, solver) keys to draw.

coadapt.rollout

rollout.py — the one loop that runs a whole episode: observe -> act -> step -> communicate.

Each round, for every agent i at once (vmapped over the agent axis):
  1. observe : obs_i = own measurement (from env)  ++  delayed parent message (from network)
  2. act     : actions, dist_params, logp = agents act on their obs
  3. step    : env turns actions into the next world + the shared team reward
  4. commun. : pack this round's messages and push them into the network's delay buffer

We `scan` over time (one compiled loop) and `vmap` over seeds (many episodes in parallel).
`horizon` and `n_seeds` are arguments, so any run can use an arbitrary episode length and an
arbitrary number of seeds — this is what makes the harness scale on a GPU.

A Traj bundles the per-round records we need for learning and for metrics:
    obs   [T, N, obs_dim]   actions [T, N, ...]   logp [T, N]
    reward[T]               per_actor_error [T, N]   dist [T, N] (hops to source, constant)

obs_dim(env, net)

The width of an agent's observation vector = own measurement + one parent message.
Also checks the 1:1 binding: the env's actor count must equal the network's node count
(each solver then builds exactly this many independent policies).

observation(env_meas, channel_msg)

Assemble each agent's observation: own measurement ++ delayed parent message.
This is the ONLY function that decides what a policy sees, so partial observability is
localized here and in network.ping_channel.

rollout(env, net, agents, key, horizon=None)

Run ONE episode. Returns a Traj (dict of [T, ...] arrays).

rollout_batch(env, net, agents, key, n_seeds, horizon=None)

Run `n_seeds` independent episodes in parallel (vmap over seeds).
Returns a Traj with a leading seed axis: [S, T, ...] (all fields are batched over seeds,
including the per-round dist -> [S, N]).

coadapt.solvers

Solvers: things that learn a policy for the N agents, and reference bounds.

Every learning solver is a MODULE exposing the same two entry points, so the runner and the
comparison script can treat them uniformly and a new baseline plugs in by adding one line
to SOLVERS below:

    solve(env, net, key, *, n_seeds, horizon, n_iters, **hp)
        -> {"return_curve": [...], "final_return": float, ...}
    default_hparams(env_name) -> dict     # sensible per-env hyperparameters

Learning solvers (each an independent learner PER AGENT, no parameter sharing):
    reinforce  — vanilla policy gradient (baseline of baselines)
    ppo        — clipped surrogate + per-agent critic + GAE
    grpo       — critic-free, group-relative advantage over the seed batch
    trpo       — KL-constrained natural-gradient step (conjugate gradient + line search)
    recurrent  — PPO whose policy carries its OWN GRU hidden state (a belief over stale
                 messages); self-contained, does not touch agent.py

Reference bounds (not learners): random floor and oracle ceiling, in random_oracle.py.

coadapt.solvers.common

common.py — shared building blocks for the from-scratch solvers.

Kept in one place so every solver (ppo, grpo, trpo, recurrent) reads the same way and we do
not repeat the returns / GAE / Adam code. Nothing here touches the environment or the agent
definition — these are pure array helpers.

actor_params_all(actor, obs)

Run every agent's actor over a batched obs [S,T,N,obs_dim] -> dist params [S,T,N,P].
vmap innermost over the agent axis, then over T, then over S.

adam_init(params)

(no docstring)

adam_step(params, grads, m, v, t, lr, b1=0.9, b2=0.999, eps=1e-08)

One Adam update over a whole pytree (per-parameter, hence per-agent).

critic_values_all(critic, obs)

Run every agent's critic over obs [S,T,N,obs_dim] -> values [S,T,N].

discounted_returns(rewards, gamma)

rewards [S, T] -> discounted-to-go G [S, T] (reverse scan over time).

final_return(curve, n_iters)

Mean per-step team reward over the last 10% of iterations (same metric for all
solvers, so the comparison table is apples-to-apples).

gae(rewards, values, gamma, lam)

Generalized Advantage Estimation.

rewards [S, T]        shared team reward
values  [S, T, N]     per-agent value estimates
returns adv [S, T, N], ret [S, T, N]  (ret = adv + values, the critic target)

standardize_per_agent(adv)

Zero-mean/unit-std per agent over the (seed, time) batch — stabilizes the gradient.

coadapt.solvers.grpo

grpo.py — Group Relative Policy Optimization, from scratch, per agent.

GRPO is PPO WITHOUT a critic. Instead of learning a value baseline, it uses the batch of
parallel episodes as a "group": an agent's advantage for an episode is its return minus the
group mean, divided by the group std. Our rollout already runs `n_seeds` episodes in
parallel, so the group is free — this makes GRPO the cheapest of the strong baselines here
(no second network to train).

Same as ppo.py otherwise: independent per-agent actors (no parameter sharing), a clipped
surrogate, several epochs per batch, and greedy evaluation.

The advantage is per-EPISODE (one number per seed per agent), broadcast across time, which is
the group-relative signal REINFORCE-with-a-mean-baseline lacks the normalization of.

default_hparams(env_name)

(no docstring)

solve(env, net, key, *, n_seeds=16, horizon=None, n_iters=200, lr=0.003, gamma=0.98, clip=0.2, epochs=4, ent_coef=0.01, hidden=64)

Train N independent GRPO agents. Returns {policy, return_curve, final_return}.

coadapt.solvers.ppo

ppo.py — from-scratch PPO, one INDEPENDENT learner per agent.

Same decentralized setup as reinforce.py (each agent has its own actor, no parameter
sharing), but with the three things that make PPO strong and stable:
  1. a per-agent CRITIC (a second MLP, reusing agent.forward) that estimates each agent's
     value, giving a low-variance advantage via GAE,
  2. a CLIPPED surrogate objective so each update cannot move the policy too far,
  3. several epochs of minibatch-free full-batch updates per collected rollout.

Everything is vmapped over the agent axis, so the N critics/actors are independent but run in
one call. We reuse agent.py's MLP for BOTH the actor (outputs action-dist params) and the
critic (outputs a single value) — the critic is just an MLP with param_size = 1. This keeps
the change confined to solvers/ (agent.py is untouched).

Evaluation metric (`final_return`) uses GREEDY actions (argmax / distribution mean) so the
reported number reflects the learned policy without exploration noise — this is what lets
PPO get close to the oracle ceiling.

default_hparams(env_name)

(no docstring)

solve(env, net, key, *, n_seeds=16, horizon=None, n_iters=200, lr=0.003, gamma=0.98, lam=0.95, clip=0.2, epochs=4, ent_coef=0.01, vf_coef=0.5, hidden=64)

Train N independent PPO agents. Returns {policy, critic, return_curve, final_return}.
return_curve[it] is the GREEDY mean per-step team reward (comparable to oracle/random).

coadapt.solvers.random_oracle

random_oracle.py — reference policies that bound every result.

  random_return : uniform-random actions. The floor any learner should beat.
  oracle_return : the best simple decentralized policy for the relay task — copy the
                  freshest thing you can about the target. On the relay worlds this is the
                  "copy your delayed parent, or use the target if you are the source" rule.
                  It is the achievable ceiling given the staleness in the channel, so it
                  shows how much of the collapse is the SETTING (not the learner).

Both return the mean per-step team reward, averaged over seeds — directly comparable to
REINFORCE's return_curve entries.

oracle_return(env, net, key, n_seeds=16, horizon=None)

Copy-the-freshest-target rule. obs = [own meas (meas_dim), channel msg (msg_dim)].
For the relay worlds: measurement[:,0]=is_source, measurement[:,1]=target if source.
The channel message's first `feature_size` slots = parent's normalized action.

random_return(env, net, key, n_seeds=16, horizon=None)

(no docstring)

coadapt.solvers.recurrent

recurrent.py — PPO with a self-contained per-agent GRU (a belief over stale messages).

This is the partial-observability baseline. Memoryless policies cannot integrate DELAYED
messages into an estimate of the moving target, so they hit a structural ceiling below the
oracle. A recurrent policy carries a hidden state through time, letting each agent build a
belief from the history of its own (stale) observations — the standard POMDP remedy.

Constraint honored: the recurrence lives ENTIRELY inside this solver. It defines its own GRU
parameters and its own rollout; it does NOT modify agent.py or the environment. The problem
setup (env, network, obs) is unchanged — only the policy that consumes the obs has memory.

Each agent has its OWN GRU + heads (no parameter sharing), all stacked on a leading [N,...]
axis and vmapped, exactly like the memoryless agents. We reuse the environment/network step
functions but run our own scan so we can thread the hidden state.

default_hparams(env_name)

(no docstring)

solve(env, net, key, *, n_seeds=16, horizon=None, n_iters=200, lr=0.003, gamma=0.98, lam=0.95, clip=0.2, epochs=4, ent_coef=0.01, vf_coef=0.5, hidden=64)

Train N independent recurrent-PPO agents. Returns {policy, return_curve, final_return}.
Actor and critic share the GRU trunk (one net per agent with two heads).

coadapt.solvers.reinforce

reinforce.py — REINFORCE written from scratch, one INDEPENDENT learner per agent.

No RL library. The only borrowed machinery is `jax.grad` for autodiff; the algorithm
(returns, baseline, loss, Adam update) is all here in plain code so the flow is readable.

Why "one independent learner per agent":
  - The reward is a shared team scalar (cooperative task), so at time t every agent sees the
    same reward r_t. We form the discounted return G_t from that shared reward.
  - But each agent has its OWN policy parameters and its OWN log-probs. Agent i's loss uses
    only agent i's log-probs and its own baseline; we take the gradient of agent i's loss
    w.r.t. agent i's parameters ONLY. There is no pooling of gradients across agents.
  - Concretely: we build a per-agent loss, then `jax.grad` w.r.t. the stacked agent pytree.
    Because agent i's loss depends only on agent i's leaves, the gradient for agent j's
    leaves comes only from agent j — the learners are independent, just computed in one
    vectorised call.

Optimizer: plain Adam, hand-coded, with per-parameter moment estimates stored alongside the
policy parameters (same pytree shape).

default_hparams(env_name)

Sensible hyperparameters for this env (falls back to the discrete defaults).

solve(env, net, key, *, n_seeds=16, horizon=None, n_iters=200, lr=0.003, gamma=0.98, hidden=32, ent_coef=0.01)

Train N independent REINFORCE agents on `env` communicating over `net`.

Returns {"policy", "return_curve", "final_return"} where return_curve[it] is the mean
per-step team reward at iteration it (so it is comparable to the random/oracle baselines).

coadapt.solvers.trpo

trpo.py — Trust Region Policy Optimization, from scratch, one INDEPENDENT learner per agent.

TRPO takes the largest policy step that stays within a KL "trust region" of the current
policy. For a single agent, each iteration solves:

    maximize   surrogate advantage
    subject to mean_KL(old || new) <= delta

via the natural gradient:  step direction = F^{-1} g (g = policy gradient, F = Fisher matrix).
We never form F; we solve `F x = g` with CONJUGATE GRADIENT using Fisher-vector products, then
a backtracking LINE SEARCH scales the step so the KL constraint holds and the surrogate improves.

STRICT DECENTRALIZATION: the entire TRPO update (gradient, CG, trust-region step-scale, and the
accept/reject line search) is defined for ONE agent on ONE agent's own trajectory, then vmapped
over the agent axis. So each agent has its own natural-gradient direction, its OWN trust-region
scale, and its OWN accept decision — no agent's update depends on another's curvature or KL.
(The earlier version summed KL/surrogate over agents and used a single global step-scale and a
joint accept, which coupled the agents; that is fixed here.)

A per-agent critic (like PPO) supplies the GAE advantage. The critic MSE is a sum of per-agent
terms, so its gradient is also independent per agent.

default_hparams(env_name)

(no docstring)

solve(env, net, key, *, n_seeds=16, horizon=None, n_iters=200, gamma=0.98, lam=0.95, delta=0.01, cg_iters=10, vf_lr=0.003, vf_epochs=4, hidden=64)

Train N independent TRPO agents. Returns {policy, critic, return_curve, final_return}.

coadapt.spaces

Action spaces and their probability distributions.

Each space knows three things so that a solver never has to special-case the action type:

    param_size          how many numbers a policy must output per agent (the head width)
    sample(params,key)  draw an action from the distribution described by `params`
    log_prob(params,a)  log-probability of action `a`
    entropy(params)     entropy of the distribution (used to encourage exploration)

`params` is always a *flat* float vector (the raw output of a per-agent MLP head). Each
space reshapes it internally, so the policy network is just `Linear(hidden -> param_size)`
regardless of the action type. All methods work with arbitrary leading batch dimensions
(e.g. [N, param_size] for N agents, or [S, T, N, param_size] for a batched rollout).

class Box(self, shape: tuple, low: float, high: float) -> None

Continuous action, diagonal Gaussian. params = [mean(D), log_std(D)] flattened.

Actions are squashed into [low, high] with a tanh so the policy cannot run off to
infinity; log_prob includes the tanh change-of-variables correction.
entropy(self, params)

(no docstring)

feature(self, action)

(no docstring)

log_prob(self, params, action)

(no docstring)

mode(self, params)

(no docstring)

sample(self, params, key)

(no docstring)

class Discrete(self, n: int) -> None

One categorical choice in {0, ..., n-1}. params = logits of length n.
entropy(self, params)

(no docstring)

feature(self, action)

(no docstring)

log_prob(self, params, action)

(no docstring)

mode(self, params)

(no docstring)

sample(self, params, key)

(no docstring)

class MultiDiscrete(self, nvec: tuple) -> None

F independent categorical choices, each in {0, ..., K-1}.

We only need the *rectangular* case (every factor has the same number of options K),
which keeps the parameters a single [..., F, K] array — required for JAX. `nvec` is
kept for readability/introspection but must be uniform.
entropy(self, params)

(no docstring)

feature(self, action)

(no docstring)

log_prob(self, params, action)

(no docstring)

mode(self, params)

(no docstring)

sample(self, params, key)

(no docstring)

coadapt.storage

storage.py — the storage harness: save/load a run as ONE structured dict file.

Every experiment is stored the same way so results are machine-readable and comparable, and
so plotting.py never has to know how a run was produced. The schema is fixed:

    {
      "config":  {...},    # what was run (env, solver, N, seeds, horizon, hyperparams)
      "metrics": {...},     # scalar summary numbers (final_return, random, oracle, ...)
      "curves":  {...},     # 1-D arrays over iterations/agents (return_curve, error_by_dist)
      "meta":    {...},     # free-form notes (git-free: timestamp passed in, versions, ...)
    }

Small scalars/lists live in a human-readable JSON; any bulky arrays go in a sidecar .npz so
the JSON stays readable. `load_run` stitches them back into one dict.

load_run(path)

Read back a run dict saved by save_run, resolving any .npz-backed arrays.

save_run(path, run)

Write a run dict to `path` (a .json). Big arrays go to `path`-with-.npz.

`run` must have the four top-level sections; missing ones default to {}.
Returns the json path written.