Skip to content

Search

Evolution-Strategies hyperparameter search for AlgorithmEntry submissions - the same algorithm navix.experiment.Experiment. run_hparam_search uses (see navix.es), but operating on any TrainingCurve-returning trainable rather than a navix Agent's own HParams specifically. This is what makes it usable for an external library's entry (e.g. rejax) that has no navix Agent/HParams at all - search_hparams only ever calls trainable(hparams, rng) and reads back a TrainingCurve, never touching whatever's inside trainable itself.

Deliberately NOT wired into Benchmark/AlgorithmEntry - Benchmark is algorithm-agnostic by design (entry.train is opaque to it; an entry "doesn't have to build a navix Agent at all", see this package's __init__.py), and hyperparameters aren't a structured, exposed field on AlgorithmEntry at all. Searching them is inherently entry-specific (what fields exist, how they're threaded into train), so it stays a tool an entry's own run.py opts into, not a Benchmark.run_env flag - see benchmarks/README.md for how a run.py uses this.

(hparams, rng) -> TrainingCurve - env_id/budget (and anything else entry.train needs) are expected to already be fixed via closure, e.g. lambda hparams, rng: entry_train_fn(hparams, env_id, budget, rng). hparams overrides only the fields named in search_hparams's own hparams_distr - everything else is whatever the closure's own defaults are.

Evolution-strategies hyperparameter search (see navix.es for the shared antithetic-sampling/probe-statistics math, and Experiment.run_hparam_search's docstring for the full algorithm description - this is the same algorithm, generalized past navix's own Agent/HParams).

Each generation: sample an antithetic population of pop_size hyperparameter sets around the current mean, call trainable on every one of them (vmapped over both the population and seeds), score each by its last-20%-mean episodic_returns (TrainingCurve. last_percent_mean, averaged over seeds), then take an ES step. The best-scoring hyperparameter set actually evaluated across every generation - not the (never directly trained) mean trajectory itself - is what's returned.

Parameters:

Name Type Description Default
trainable Trainable

(hparams, rng) -> TrainingCurve - see Trainable's own docstring for the closure convention.

 required
hparams_distr Dict[str, Distribution]

One distribution per searched field - seeds that field's starting value/scale/valid range (via navix.es. probe_hparam_field_stats), not resampled every generation. Every candidate and every ES update is clipped to that field's empirical [min, max] - without it, nothing stops the search drifting a field outside the range this distribution was ever meant to describe (e.g. a gae_lambda distribution meant to express "search within [0.8, 0.99]" wouldn't stop the search drifting past 1.0, an invalid value).

 required
seeds Tuple[int, ...]

PRNG seeds trainable is vmapped over per candidate - must have more than one, so fitness isn't just RNG luck for a single rollout.

 required
pop_size int

Population size per generation. Must be even - antithetic sampling pairs (+/-).

 8
num_generations int

Number of ES update steps.

 10
sigma float

Noise scale, in units of each field's own empirical probe std.

 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

Returns:

Type Description
Dict[str, float]

Tuple[Dict[str, float], Array]: The best-scoring hyperparameter
Array

set actually evaluated across every generation (plain Python
Tuple[Dict[str, float], Array]

floats, ready to splice into whatever config trainable's
Tuple[Dict[str, float], Array]

closure builds from), and its fitness (last-20%-mean
Tuple[Dict[str, float], Array]

episodic_returns, averaged over seeds).

Raises:

Type Description
ValueError

If seeds has one or fewer entries, or pop_size is odd.