Experiment

Experiment: trains one Agent against one Environment, across seeds, optionally logging to wandb (run), or searches its hyperparameters via Evolution Strategies (run_hparam_search).

Experiment

Trains one agent on one env for a list of seeds and collects the metrics.

run() compiles the agent's train once and runs it for every seed (each seed is an independent training run from a different PRNG key), returning the stacked final train state and a logs pytree with a leading seed axis. run_hparam_search() wraps that in an Evolution-Strategies loop over hyperparameters.

Parameters:

Name	Type	Description	Default
name	str	experiment name; used as the wandb project and as a prefix for run names.	required
agent	Agent	the agent to train (PPO, PQN, Dreamer, or a custom Agent). It already holds its own hparams.	required
env	Environment	the environment to train on. Usually the same one nx.make(env_id) returns.	required
env_id	str	the registered id of env (e.g. "Navix-Empty-5x5-v0"). Logged as metadata and used to match a run to a navix.benchmarks protocol; "" if env is not a registered environment.	''
seeds	Tuple[int, ...]	one training run per seed. Default (0,).	(0,)
group	str	optional wandb group, for aggregating related runs in the dashboard.	''

run(log_to_wandb=True, do_log=None)

Default function to run the experiment. This function compiles the training function, trains the agent, and logs the results.

Two strategies exist for looking at the results, and they trade off against each other:

• log_to_wandb=True (the default) streams metrics to Weights & Biases as training progresses. This is the slow path - real network I/O, roughly linear in the number of seeds.
• log_to_wandb=False skips wandb entirely; logs (this method's second return value) is returned either way, so with wandb off you get it back much faster, with no network calls at all. Pair it with navix.benchmarks.plotting to get a local matplotlib dashboard from logs instead of a wandb one. See issue #60.

Parameters:

Name	Type	Description	Default
log_to_wandb	bool	Whether to log the results to wandb.	True
do_log	bool	Deprecated alias for log_to_wandb.	None

Returns:

Name	Type	Description
Tuple		A tuple containing the final training state and the logs.

run_hparam_search(hparams_distr, pop_size, num_generations=10, sigma=1.0, solver=None, n_probe=256, log_to_wandb=True)

Evolution-strategies hyperparameter search, adapted from OpenAI-ES (Salimans et al., 2017 - https://arxiv.org/abs/1703.03864) the way https://github.com/ESHyperscale/HyperscaleES's open_es.py applies it to neural network weights, here applied to a hyperparameter vector instead.

Each generation: sample an antithetic (mirrored +/-) population of pop_size hyperparameter sets around the current mean, train all of them in one fused jax.jit(jax.vmap(...)) call (the same shape Experiment.run itself uses), score each by its last-20%-mean agent/episode/returns (navix.agents.agent.derive_episodic_metrics, averaged over self.seeds), then take an ES step: z-score the fitnesses, estimate a gradient from fitness-weighted noise, and update the mean via solver. The best-scoring hyperparameter set actually evaluated across every generation - not the (never directly trained) mean trajectory itself - is what's returned.

Every searched field spans a different natural scale (lr ~1e-4, gae_lambda ~0.95, ...), so sigma is relative, not absolute: an empirical probe of n_probe samples from each field's own hparams_distr distribution gives that field's starting value (the probe's mean) and natural scale (the probe's std) - sigma is then how many of those per-field stds each generation's noise spans. This also sidesteps relying on a distribution's .mean()/ .stddev(), which a distribution like examples/hparam_search.py's CategoricalUniform (its .sample() maps a Categorical's sampled index through a domain list, but doesn't override .mean()/ .stddev() to match) would silently get wrong.

Parameters:

Name	Type	Description	Default
hparams_distr	Dict[str, Distribution]	One distribution per searched field. Keys must name a field on self.agent.hparams with pytree_node=True (navix's continuous float hparams - lr, clip_eps, gae_lambda, ... - not budget/num_envs/etc., which stay structurally unsearchable). Each distribution seeds that field's starting value/scale (see above) - it does not keep resampling every generation.	required
pop_size	int	Population size per generation. Must be even - each generation samples pop_size // 2 noise vectors and mirrors them (antithetic sampling), the same variance-reduction trick open_es.py uses.	required
num_generations	int	Number of ES update steps.	10
sigma	float	Noise scale, in units of each field's own empirical probe std (see above).	1.0
solver	GradientTransformation	The ES mean update rule. Defaults to optax.sgd(0.1).	None
n_probe	int	Samples drawn from each field's distribution to estimate that field's starting value and scale.	256
log_to_wandb	bool	Whether to log per-generation fitness stats plus the final best candidate's training curve to wandb.	True

Returns:

Type	Description
HParams	Tuple[HParams, Array]: The best-scoring hyperparameter set
Array	actually evaluated across every generation, and its fitness
Tuple[HParams, Array]	(last-20%-mean agent/episode/returns, averaged over self.seeds).

Raises:

Type	Description
ValueError	If pop_size is odd, or hparams_distr names a pytree_node=False field.

build_search_set(base_hparams, candidates, pop_size)

Batches pop_size individually-replace'd copies of base_hparams (one per population member's candidates) into a single HParams pytree with a new leading population axis - ready for jax.vmap.

Parameters:

Name	Type	Description	Default
base_hparams	HParams	Every other (non-searched) field's value.	required
candidates	Dict[str, Array]	This generation's per-field candidate values, each shaped (pop_size,) (see navix.es.sample_antithetic_candidates).	required
pop_size	int	Population size.	required

Returns:

Name	Type	Description
HParams	HParams	Every searched field's leaves shaped (pop_size, ...).

hparam_search_fitness(logs)

One scalar fitness per population member, from logs (as returned by a run_hparam_search generation's search_fn call): last-20%-mean agent/episode/returns (navix.agents.agent. derive_episodic_metrics), averaged over the seed axis.

Parameters:

Name	Type	Description	Default
logs	Dict[str, Array]	Shape (pop_size, num_seeds, num_updates, num_steps, num_envs) for agent/train/done_mask/ agent/train/returns/agent/train/lengths.	required

Returns:

Name	Type	Description
Array	Array	Shape (pop_size,).