Robust evaluation

Robust evaluation of a policy over the distribution an environment samples from.

evaluate runs n independent episodes of env under a fixed policy and reports the physical final performance (mean, lower-tail CVaR, minimum) with bootstrap confidence intervals where they are validated. The evaluation distribution is whatever the environment carries: wrap the task in a parameter randomiser, a control-error wrapper or an observation model first; there is no separate distribution argument. Evaluate a policy on a distribution other than its training distribution by building the evaluation environment separately and passing a key that the training stream did not use.

What is reported

metric: the physical final performance, read from the exact env.info(final_state): "fidelity" for gate and state-transfer tasks, the reward quantity ("energy" or "ergotropy") for battery tasks. episode_return (the training reward summed over the episode, a gain Q(final) - Q(initial)) is reported in separate fields; aggregation is always over the metric. The metric is higher is better, so the CVaR here is the mean of the lowest ceil(0.1 n) episodes, the opposite tail of a loss-based CVaR.
mean, cvar (lower 10 percent tail mean) and min. The CVaR estimator has an O(1/n) bias (for U(0, 1) its expectation is exactly (k + 1) / (2 (n + 1)), k = ceil(0.1 n), against the population value 0.05). The minimum is a descriptive statistic that falls as n grows: never compare minima of different n.
Intervals: non-parametric percentile bootstrap over episodes (B = 2000, 2.5 and 97.5 percentiles). The mean interval is provided for n >= 256 and the CVaR interval for n >= 1024; below that only point estimates are returned, with a warning. In a Monte Carlo check (1000 independent repetitions each) the empirical coverage at n = 1024 was inside [0.92, 0.98] for U(0, 1) and for a gate-fidelity-like distribution, but the CVaR interval undercovers for U(0, 1) (0.932, about 2.6 Monte Carlo standard errors below the nominal 0.95), so it is approximate and somewhat narrow. That is the evidence for those two distributions, not a coverage guarantee for any distribution: heavy-tailed, discrete or boundary-heavy performance distributions (for example a fidelity piled up at its upper bound 1) are separate validation targets.
The intervals describe the sampling variation of the evaluation episodes for the given policy. Variation across independently trained policies (seeds) is a different source; report it separately and do not merge the two into one interval.

Random numbers: key is required. The episode keys are keys.episode_key(fold_in(key, 0), i) for i = 0, ..., n - 1, the bootstrap key is fold_in(key, 1) and a stochastic policy receives keys.policy_key(episode_key, t) at decision t = 0, 1, .... Use a key that does not belong to the training stream. The outcome of an episode does not depend on the batch it is simulated in, up to the last bits of floating-point rounding (XLA changes the operation order between batch sizes).

`CVAR_LEVEL = 0.1` `module-attribute`

Fraction of the episodes in the lower tail averaged by :func:cvar.

`N_BOOT = 2000` `module-attribute`

Bootstrap replicates B of the percentile intervals.

`Evaluation` `dataclass`

Result of :func:evaluate.

Attributes:

Name	Type	Description
`metric`	`str`	name of the aggregated performance metric.
`n`	`int`	number of episodes.
`mean`	`float`	mean of the metric.
`cvar`	`float`	mean of the lowest `ceil(0.1 n)` metric values (O(1/n) biased estimator).
`min`	`float`	smallest metric value (descriptive; falls with `n`).
`mean_interval`	`tuple[float, float] \| None`	95 percent bootstrap interval of the mean, `None` for `n < 256`.
`cvar_interval`	`tuple[float, float] \| None`	95 percent bootstrap interval of the CVaR, `None` for `n < 1024`.
`return_mean`	`float`	mean episode return (training reward), reported apart from the metric.
`return_min`	`float`	smallest episode return.
`nominal`	`float \| None`	metric of the same policy on `nominal_env` (mean of 2 episodes of an environment that is assumed deterministic), `None` if not requested.
`nominal_return`	`float \| None`	its episode return.
`values`	`ndarray`	per-episode metric, shape `(n,)`, in episode order.
`returns`	`ndarray`	per-episode returns, shape `(n,)`.

`as_dict()`

Plain-dict summary (floats and lists, without the per-episode arrays).

`Stability` `dataclass`

Result of :func:stability: the evaluation at several sample sizes.

Attributes:

Name	Type	Description
`results`	`dict[int, Evaluation]`	`{n: Evaluation}`; the episodes of the smaller sizes are the first episodes of the larger ones (same episode keys).
`mean_spread`	`float`	`max - min` of the mean over the sizes.
`cvar_spread`	`float`	`max - min` of the CVaR over the sizes (the minimum is not compared across sizes: it is a descriptive statistic that falls with `n`).

`bootstrap_intervals(values, key, *, mean=True, tail=True, n_boot=N_BOOT, level=0.95)`

Percentile bootstrap intervals of the mean and of the CVaR of values.

Resamples the n episode values with replacement n_boot times (both statistics use the same resamples) and returns the (1 - level) / 2 and (1 + level) / 2 percentiles of each statistic, keyed "mean" and "cvar". The resampling indices come from a NumPy generator seeded with the data of key, so the result is a function of (values, key) only.

`cvar(values, axis=-1)`

Lower-tail CVaR: mean of the lowest ceil(0.1 n) entries along axis.

Higher is better, so this is the average of the worst 10 percent of the episodes. Ties are handled exactly (the tail is a multiset of values). Biased by O(1/n).

`evaluate(env, policy, key, n=1024, metric=None, nominal_env=None, *, batch_size=None)`

Evaluate policy on n independent episodes of env (see the module docs).

Parameters:

Name	Type	Description	Default
`env`	`Environment`	environment to evaluate in; carries the evaluation distribution (wrappers).	required
`policy`	`Policy`	an open-loop action array of shape `(horizon, action_dim)` or a JAX-traceable function `(observation, key) -> action` (stochastic policies use `key`, a fresh key `keys.policy_key(episode_key, t)` at every decision `t`).	required
`key`	`Array`	explicit evaluation key (required); see the module docs for the derivation.	required
`n`	`int`	number of episodes (default 1024; `n >= 256` for a mean interval, `n >= 1024` for a CVaR interval).	`1024`
`metric`	`Metric`	an `env.info` key, a function of the final state returning a scalar, or `None` for :func:`performance_metric`.	`None`
`nominal_env`	`Environment \| None`	the same task without randomisation. The policy is also run on it for 2 episodes (it is assumed deterministic; a warning is issued if the two episodes differ) and the result is reported as `Evaluation.nominal`, so that robustness bought at the expense of nominal performance is visible.	`None`
`batch_size`	`int \| None`	simulate the episodes in chunks of this size to bound memory (default: all at once).	`None`

Returns:

Type	Description
`Evaluation`	class:`Evaluation`.

Raises:

Type	Description
`ValueError`	no key, a non-positive `n`, an open-loop array of the wrong shape, or an environment without a default metric and no `metric`.

`performance_metric(env)`

Name of the physical final performance of the task behind env.

Inferred from the unwrapped environment: GateControl (average gate fidelity, leakage included) and StateTransfer (target-state fidelity) give "fidelity"; BatteryCharging gives its reward quantity, "energy" or "ergotropy". Any other environment raises ValueError: pass metric= to :func:evaluate explicitly.

`stability(env, policy, key, sizes=DEFAULT_SIZES, metric=None, *, batch_size=None)`

Mean and CVaR at several sample sizes, to show whether they have settled.

One run of max(sizes) episodes is simulated; the result for a smaller size uses the first n of those episodes (the same episode keys as evaluate(..., n=n)), so the sizes are nested and share one bootstrap key. Intervals follow the rules of :func:evaluate for each size (no interval below its minimum n); no warning is issued because showing the small sizes is the purpose of this report.

`tail_count(n, level=CVAR_LEVEL)`

Number k = ceil(level * n) of episodes in the lower tail (exact for 0.1).

Key tree

The key tree of randomised environments.

Every random draw of an episode is derived from the episode key alone, through a fixed tree of jax.random.fold_in calls, so that changing how one leaf is used (how many draws, which distribution, on or off) never changes the samples of another leaf::

run_key
episode_key[i] = fold_in(run_key, i)
    0 parameters             fold_in(episode_key, 0)
        one key per name     fold_in(., crc32(name))
    1 control errors         fold_in(episode_key, 1)
        0 calibration        fold_in(., 0)
            0 gain error     fold_in(., 0)
            1 offset         fold_in(., 1)
        1 step noise         fold_in(., 1)
            step k           fold_in(., k)
    2 measurement            fold_in(episode_key, 2)
        observation t        fold_in(., t)   (t = 0 after reset, t = k after step k)
            observable j     fold_in(., j)
    3 policy (evaluation)    fold_in(episode_key, 3)
        decision t           fold_in(., t)

Keys are a function of the episode key and the step index only, so restarts and parallel runs do not depend on call order. Statistical independence between leaves is not proved by this construction; the tests check sample distributions and correlations empirically.

`calibration_key(ep_key, which)`

Key of an episode-fixed control error: GAIN or OFFSET.

`check_name_tags(names)`

Raise ValueError if two different names share a tag.

`episode_key(run_key, i)`

Key of the i-th episode of a run (the rule used by vector envs and evaluation).

`measurement_key(ep_key, t, j)`

Key of the finite-shot estimate of observable j in observation t.

`name_tag(name)`

Stable 31-bit tag of a parameter name (CRC32, so it does not depend on the order).

`parameter_key(ep_key, name)`

Key of the model parameter name (fixed for the episode).

`policy_key(ep_key, t)`

Key handed to a stochastic policy at decision t (evaluation only).

`step_noise_key(ep_key, k)`

Key of the control noise of step k.

Robust evaluation

What is reported

CVAR_LEVEL = 0.1 module-attribute

N_BOOT = 2000 module-attribute

Evaluation dataclass

as_dict()

Stability dataclass

bootstrap_intervals(values, key, *, mean=True, tail=True, n_boot=N_BOOT, level=0.95)

cvar(values, axis=-1)

evaluate(env, policy, key, n=1024, metric=None, nominal_env=None, *, batch_size=None)

performance_metric(env)

stability(env, policy, key, sizes=DEFAULT_SIZES, metric=None, *, batch_size=None)

tail_count(n, level=CVAR_LEVEL)