Es

Evolution-Strategies primitives shared by every navix hyperparameter search: Experiment.run_hparam_search (navix/experiment.py, searches a navix Agent's own HParams) and navix.benchmarks.search. search_hparams (searches an arbitrary AlgorithmEntry-shaped trainable's hyperparameters, e.g. an external library like rejax). Both follow the same OpenAI-ES (Salimans et al., 2017 - https://arxiv.org/ abs/1703.03864) shape - antithetic Gaussian sampling around a per-field- scaled mean - only how a generation's population actually gets trained and scored differs, which is why that part isn't shared here.

LogUniform

Bases: Uniform

Log-uniform over [low, high] (both > 0): log(sample) is uniform in [log(low), log(high)), so every order of magnitude in the range gets equal sampling weight - the right choice for a field like a learning rate, where the right order of magnitude matters far more than the right value within one. A plain distrax. Uniform(low, high) over several orders of magnitude does not do this: Uniform(1e-5, 1e-2) puts ~90% of its mass above 1e-3, starving the smaller end where a value like PPO's own 2.5e-4 default typically lives - confirmed in practice, searching navix's own PPO with a plain Uniform(1e-5, 1e-2) for lr found values clustered at 1e-3-1e-2 across every environment, nowhere near the known-good region.

__init__(low, high)

Args: low (float): lower bound, > 0. high (float): upper bound, > low.

sample(*, seed, sample_shape=())

Draws a log-uniform sample in [low, high].

Parameters:

Name	Type	Description	Default
seed		a jax.random PRNG key.	required
sample_shape		extra leading shape for a batch of draws.	()

Returns:

Name	Type	Description
Array		the sample(s), always positive.

sample_n(rng, n)

n log-uniform samples, shape (n, ...).

Parameters:

Name	Type	Description	Default
rng		a jax.random PRNG key.	required
n	int	how many samples.	required

probe_hparam_field_stats(hparams_distr, n_probe, key)

Empirically estimates each searched field's starting value, scale, and valid range from n_probe samples of its own distribution - more robust than relying on a distribution's .mean()/.stddev(), which silently gets the wrong answer for 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).

Parameters:

Name	Type	Description	Default
hparams_distr	Dict[str, Distribution]	One distribution per searched field.	required
n_probe	int	Samples drawn per field.	required
key	Array	PRNG key, split once per field.	required

Returns:

Type	Description
Dict[str, Array]	Tuple[Dict, Dict, Dict, Dict]: (theta, scale, lo, hi) - each
Dict[str, Array]	field's probe mean, probe std (floored to avoid a degenerate
Dict[str, Array]	zero-sigma field), and probe min/max. lo/hi bound every
Dict[str, Array]	candidate sample_antithetic_candidates produces and every ES
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	update step - without them, nothing stops the search's random
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	walk drifting a field outside the range hparams_distr was
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	ever meant to describe (confirmed in practice: an unclipped
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	search drifted gae_lambda - only ever valid in [0, 1] - to
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	1.005, and separately drifted a learning_rate down to
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	exactly 0.0 and got stuck there, since a learning rate of
Tuple[Dict[str, Array], Dict[str, Array], Dict[str, Array], Dict[str, Array]]	zero produces no fitness signal to climb back out with).

sample_antithetic_candidates(theta, scale, lo, hi, pop_size, sigma, key)

One ES generation's population: pop_size // 2 i.i.d. standard- normal noise vectors per field, mirrored (antithetic sampling) to fill the rest of the population, then theta + sigma * scale * noise per field, clipped to [lo[k], hi[k]] (see probe_hparam_field_stats).

Parameters:

Name	Type	Description	Default
theta	Dict[str, Array]	Current per-field mean.	required
scale	Dict[str, Array]	Per-field noise scale (see probe_hparam_field_stats).	required
lo	Dict[str, Array]	Per-field lower bound.	required
hi	Dict[str, Array]	Per-field upper bound.	required
pop_size	int	Population size. Must be even.	required
sigma	float	Noise scale, in units of scale.	required
key	Array	PRNG key, split once per field.	required

Returns:

Type	Description
Dict[str, Array]	Tuple[Dict, Dict]: (noise, candidates), each `Dict[str,
Dict[str, Array]	Array]shaped(pop_size,)per field -noise` is what the ES
Tuple[Dict[str, Array], Dict[str, Array]]	gradient estimate is computed from, candidates is what
Tuple[Dict[str, Array], Dict[str, Array]]	actually gets trained.