Skip to content

Benchmark

TrainingCurve (what AlgorithmEntry.train returns - purely what's computable inside a jax.jit trace), AlgorithmEntry (the provenance and hardware metadata wrapper for one algorithm being scored, issue #130) plus its CostAnalysis, BenchmarkResult (one entry's scored run - a TrainingCurve plus everything only an external, un-jitted wrapper can measure: wall-clock time and cost), and Benchmark (the abstract protocol base every preset, e.g. navix.benchmarks.scratch.FromScratchBenchmark, subclasses). See navix.benchmarks (this package's __init__.py) for the full design.

One algorithm to score against a Benchmark.

To submit an algorithm: subclass AlgorithmEntry, override train to build whatever model env_id needs and train it and return a TrainingCurve - the only requirement - then construct an instance with the provenance fields below (see benchmarks/README.md). The hardware fields (gpu_type through jaxlib_version) are auto-detected in __post_init__, not constructor arguments. __post_init__ also checks that train returns a TrainingCurve with the right shape, so a malformed train fails at construction time, not partway through a real benchmark run.

Attributes:

Name Type Description
name str

Algorithm name, e.g. "PPO".
author str

This implementation's author (GitHub handle), not the paper's. Validated in __post_init__.
paper_url str

Link to the paper the algorithm is from.
navix_commit_url str

Link to the navix commit this result was produced against (issue #130's navix.sha, as a URL). Validated in __post_init__.
algorithm_commit_url str

Link to the algorithm implementation's own commit (issue #130's agent.sha, as a URL) - same commit as navix_commit_url for a navix-shipped agent, a different repo's commit for an external one. Validated in __post_init__.
gpu_type Optional[str]

The GPU model JAX runs on, or None if JAX isn't running on a GPU.
cpu_type str

The CPU's model name.
ram_bytes int

Total system RAM, in bytes.
cuda_version Optional[str]

The CUDA version jaxlib runs on, or None if JAX isn't running on a GPU.
cudnn_version Optional[str]

The cuDNN version jaxlib runs on, or None if JAX isn't running on a GPU.
jax_version str

jax.__version__.
jaxlib_version str

jaxlib.__version__.

Compiles self.train(env_id, budget, ...) and reads its FLOPs/memory/compile-time. Always seed 0, since cost is shape-driven, not value-driven.

Necessarily includes whatever env interaction train does internally. In practice this lands close to "one update's cost", not the whole run's: every navix-shipped agent's train is init + jax.lax.scan(self.update, ..., length=num_updates), and XLA's cost_analysis() on a compiled scan reports one iteration's cost, not length copies. An agent whose train isn't scan-shaped will report a different figure here.

Parameters:

Name Type Description Default
env_id str

The environment to build the model for.

 required
budget int

Training budget, passed to self.train.

 required

Returns:

Name Type Description
CostAnalysis CostAnalysis

FLOPs, peak memory, and compile time for one
CostAnalysis

self.train(env_id, budget, ...) call.

Builds a fresh, env-shaped model for env_id and trains it at budget. The one method every submission must override - building the model is inherently algorithm-specific, so there's no protocol-agnostic default.

Runs inside a jax.jit/jax.vmap trace (see Benchmark. run_env), so only genuinely jittable work belongs here - wall-clock timing and cost are measured separately, from outside any trace (see BenchmarkResult, cost_analysis).

Parameters:

Name Type Description Default
env_id str

The environment to train on.

 required
budget int

Training budget - whatever the running Benchmark protocol passed to Benchmark.run_env. Use it to build your hparams (e.g. PPOHparams(budget=budget)) if your algorithm's training length should respect it.

 required
rng Array

The PRNG key to train with.

 required

Returns:

Name Type Description
TrainingCurve TrainingCurve

episodic_returns/lengths filled in;
TrainingCurve

diagnostics optionally, with whatever per-update values
TrainingCurve

help debug this algorithm.

Raises:

Type Description
NotImplementedError

Always, on AlgorithmEntry itself - must be overridden.

Checks self.train returns a TrainingCurve with the right shape, without running any real training - jax.eval_shape traces self.train for its output structure only, against one representative registered environment. A single (un-vmapped) call, so episodic_returns/lengths/diagnostics' values must all be rank 1 (one point per update).

Raises:

Type Description
TypeError

If self.train's output isn't a TrainingCurve.
AssertionError

If episodic_returns/lengths/ diagnostics' values aren't rank 1 (chex.assert_rank).

Bases: PyTreeNode

An experimental protocol - the fixed set of choices that make two algorithms' numbers comparable: which environments, what frame budget, how many seeds, and how a run is scored and summarised. It says nothing about how an algorithm is implemented, only how it is measured, so any AlgorithmEntry can be run against it and against the literature.

Use an instance, e.g. Navix1M().run(entry); the protocol's name is type(self).__name__. A concrete protocol (e.g. FromScratchBenchmark, behind Navix1M / Navix100K) supplies run / summary / details; submit_entry is shared - it writes out whatever those produced, identically for every protocol.

Reduces a BenchmarkResult from run into per-row diagnostics about this benchmark run - the same kind of columns summary aggregates, but one row per whatever this protocol's leading axis represents (e.g. one row per environment for FromScratchBenchmark), instead of a single further-aggregated row. Diagnostics about the benchmark run itself, not a leaderboard's per-algorithm click-through page - what that page shows depends on the algorithm, not the benchmark protocol.

A concrete override should include whatever labels row identity (e.g. env_ids) as one of its own returned entries, if this protocol has a meaningful one - submit_entry writes out exactly what this returns and nothing more. Values aren't required to be Array (e.g. env_ids is a Tuple[str, ...], not every metric need be jax-typed) - unlike summary, this isn't reduced to a uniform type.

Parameters:

Name Type Description Default
results BenchmarkResult

This protocol's run output.

 required

Returns:

Type Description
Dict[str, Any]

Dict[str, Any]: One row per unit of this protocol's
Dict[str, Any]

leading axis - which columns exist, and what that axis is,
Dict[str, Any]

is protocol-specific.

Raises:

Type Description
NotImplementedError

Always, on Benchmark itself - must be overridden by a concrete protocol.

self.details(results) as a local, offline bar-chart figure, one panel per numeric metric, one bar per row (e.g. one bar per environment for FromScratchBenchmark) - mean plus a std-dev error bar over whatever trailing axis details keeps raw (e.g. FromScratchBenchmark keeps every seed's own value, unlike summary's already-averaged numbers - see FromScratchBenchmark.details's docstring).

Parameters:

Name Type Description Default
results BenchmarkResult

This protocol's run output.

 required

Returns:

Type Description

matplotlib.figure.Figure: One panel per numeric metric.

results.curve's raw training curves as a local, offline figure, one panel per curve (episodic_returns, length, any curve.diagnostics entries) - mean line plus a min-max band over self.seeds, same convention as navix.benchmarks. plotting.plot_metric. Same curves submit_entry writes (resampled) into diagnostics.npz, but at full resolution and without needing a file round-trip.

A TrainingCurve doesn't carry absolute frame counts (unlike the raw logs pytree plot_metric plots), so the x-axis is training progress as a 0-100% fraction of however many points the curve has, not a frame count.

Parameters:

Name Type Description Default
results BenchmarkResult

This protocol's run output.

 required

Returns:

Type Description

matplotlib.figure.Figure: One panel per curve.

self.summary(results) as a local, offline metric/value table figure. Independent of whatever charts the online leaderboard renders from the same summary.json.

A bar chart would be misleading here: summary's metrics live on wildly different scales in the same dict (episodic returns in [0, 1] next to flops in the hundreds of millions), so a table keeps every value legible without implying they're comparable.

Parameters:

Name Type Description Default
results BenchmarkResult

This protocol's run output.

 required

Returns:

Type Description

matplotlib.figure.Figure: The table figure.

Trains entry under this protocol.

Parameters:

Name Type Description Default
entry AlgorithmEntry

The algorithm to score.

 required

Returns:

Name Type Description
BenchmarkResult BenchmarkResult

However many things this protocol
BenchmarkResult

measures per run, stacked along a leading axis - which
BenchmarkResult

axis, and what it represents (e.g. one row per environment
BenchmarkResult

for FromScratchBenchmark), is protocol-specific.

Raises:

Type Description
NotImplementedError

Always, on Benchmark itself - must be overridden by a concrete protocol.

Trains entry on env_id at budget, vmapped over self.seeds, timing the run (compile time excluded) and reading its cost.

Parameters:

Name Type Description Default
entry AlgorithmEntry

The algorithm to train.

 required
env_id str

The environment to train on.

 required
budget int

Training budget, passed to entry.train/ entry.cost_analysis.

 required

Returns:

Name Type Description
BenchmarkResult BenchmarkResult

curve is entry.train's output, one
BenchmarkResult

per seed stacked along a new leading axis. wall_time/
BenchmarkResult

fps/cost are each a single scalar for the whole
BenchmarkResult

vmapped call - seeds train together in one fused
BenchmarkResult

computation, so there's no meaningful per-seed timing
BenchmarkResult

breakdown.

submit_entry(entry, results, max_points=50, subdir='')

Writes one self.run(entry)'s output into the directory of whichever script called this - the same convention every submission's run.py already follows for config.yml/requirements.txt, so a submission's results end up right alongside them.

Writes three files:

  • summary.json: entry's provenance/hardware fields plus self.summary(results) - the leaderboard's table row.
  • details.json: self.details(results) - per-row diagnostics about this run.
  • diagnostics.npz: results itself. curve.episodic_returns/ curve.lengths/curve.diagnostics' values (written as benchmark/episode/returns/benchmark/episode/length, plus curve.diagnostics' own keys unchanged - see curve_diagnostics) are resampled to exactly max_points evenly-spaced points along their trailing axis, so every submission's curve fields end up the same fixed shape regardless of how many updates actually ran; benchmark/costs/* (wall_time/fps/cost.*) are already scalars and are written as-is.

Parameters:

Name Type Description Default
entry AlgorithmEntry

The algorithm that produced results.

 required
results BenchmarkResult

This protocol's run output.

 required
max_points int

Number of points curve.episodic_returns/ curve.lengths/curve.diagnostics' values are resampled to in diagnostics.npz.

 50
subdir str

If non-empty, writes into a subdirectory of the caller's own directory instead of the directory itself (created if missing). For a run.py that scores the same algorithm under more than one configuration it doesn't expose as a structured AlgorithmEntry field (e.g. observation type - see navix.benchmarks.search's module docstring for the same reasoning applied to hyperparameters: what's configurable is inherently entry-specific, so Benchmark stays agnostic to it) - one submit_entry call per configuration, each into its own subdir, instead of colliding on one shared summary.json.

 ''

Reduces a BenchmarkResult from run into a leaderboard's table row for one algorithm entry.

Parameters:

Name Type Description Default
results BenchmarkResult

This protocol's run output.

 required

Returns:

Type Description
Dict[str, Array]

Dict[str, Array]: The named columns a leaderboard's table
Dict[str, Array]

row shows for this entry - which columns exist is
Dict[str, Array]

protocol-specific.

Raises:

Type Description
NotImplementedError

Always, on Benchmark itself - must be overridden by a concrete protocol.

Bases: PyTreeNode

One AlgorithmEntry's scored run under a Benchmark protocol - a TrainingCurve plus everything only an external, un-jitted wrapper can measure. Built by Benchmark.run_env in one shot, from three independently-measured pieces.

Attributes:

Name Type Description
curve TrainingCurve

AlgorithmEntry.train's output.
wall_time Array

Real wall-clock time to execute the already-compiled AlgorithmEntry.train (all of Benchmark.seeds, vmapped together) - timed and jax.block_until_ready'd from outside any jax.jit trace, unlike anything train could measure about itself. Excludes compile time (see cost.compile_time_seconds). Scalar.
fps Array

Training throughput: budget / wall_time. Scalar. Comparable only across results measured on the same hardware.
cost CostAnalysis

From AlgorithmEntry.cost_analysis.

Bases: PyTreeNode

The cost of one AlgorithmEntry.train call, as measured by AlgorithmEntry.cost_analysis.

Attributes:

Name Type Description
flops float

FLOPs, from compiled.cost_analysis().
memory_bytes float

Peak memory proxy (argument + temp + output size), from compiled.memory_analysis().
compile_time_seconds float

Wall-clock time to compile. Hardware/XLA-version-sensitive.

Bases: PyTreeNode

One AlgorithmEntry.train call's measurements - purely what's computable from inside a jax.jit trace, from the real (state, action, reward, done) interaction stream. No wall-clock timing, no cost - time.time() inside a jitted function only ever fires at trace time, not per call, so neither can be measured here; see BenchmarkResult for those.

Every field is a per-update curve - shape (num_updates,) for a single training curve (checked by AlgorithmEntry. validate_train_contract), so last_percent_mean/ last_percent_variance/convergence_rate reduce all of them the same way, uniformly, no exceptions.

Attributes:

Name Type Description
episodic_returns Array

Episodic return, masked-mean over completed episodes only.
lengths Array

Episode length, masked-mean over completed episodes only.
diagnostics Dict[str, Array]

Free-form per-update diagnostic curves (e.g. {"loss": ..., "lr": ...}) - whatever helps debug the algorithm. Each value shape (num_updates,). Empty by default.

Reduces every field to a normalized area under its curve: the ratio of the curve's overall mean to its last_percent_mean. Near 1 means the curve was close to its own asymptote for most of training (fast convergence); near 0 means most of training was spent far from it (slow).

Not bounded to [0, 1] - and shouldn't be clipped to look like it is. That range only holds when the curve improves monotonically toward its own tail. Whenever the tail is worse than the training-long average - policy collapse, instability, catastrophic forgetting, all real and fairly common RL failure modes - overall / target correctly exceeds 1, in principle arbitrarily far (confirmed in practice: a seed that learned real signal early, then collapsed to near-zero returns by the final 20% of training, produced 77 here - correctly flagging that specific seed as "learned, then collapsed" rather than "steady, healthy convergence", which a clipped value couldn't distinguish). A negative value is similarly meaningful for a negative-reward task, not an error.

Parameters:

Name Type Description Default
percent float

Percentage of the trailing axis that defines the asymptote (see last_percent_mean).

 20

Returns:

Name Type Description
TrainingCurve TrainingCurve

A copy with every field reduced along its
TrainingCurve

trailing axis.

Reduces every field to its mean over the last percent% of its trailing axis.

Parameters:

Name Type Description Default
percent float

Percentage of the trailing axis to average over.

 20

Returns:

Name Type Description
TrainingCurve TrainingCurve

A copy with every field reduced along its
TrainingCurve

trailing axis.

Reduces every field to its variance over the last percent% of its trailing axis - how much an already-converged curve still fluctuates update-to-update, not variance across seeds. The training-curve analogue of the per-update variance in Rowland, Dabney & Munos, "Adaptive Trade-Offs in Off-Policy Learning" (https://arxiv.org/abs/1910.07478, Definition 1.4).

Parameters:

Name Type Description Default
percent float

Percentage of the trailing axis to compute the variance over.

 20

Returns:

Name Type Description
TrainingCurve TrainingCurve

A copy with every field reduced along its
TrainingCurve

trailing axis.

result.curve's fields flattened into the plain {name: array} shape both Benchmark.submit_entry (writing diagnostics.npz) and Benchmark.plot_diagnostics (plotting the same curves in-memory) need: benchmark/episode/returns/benchmark/episode/ length plus curve.diagnostics's own keys, unchanged - curve. diagnostics is free-form (see TrainingCurve's docstring), so this doesn't impose or assume any prefix on it. navix's own agents happen to already key it agent/diagnostics/<name> (Agent. train's own contract, preserved end to end through benchmarks/ */*/run.py's TrainingCurve.diagnostics construction), but a submitted entry that isn't a navix agent owes no such convention.

Parameters:

Name Type Description Default
result BenchmarkResult

Already jax.device_get'd by the caller.

 required
resample callable

If given, applied to every curve array (e.g. submit_entry's fixed-max_points resampling). None keeps each curve at its full, native resolution.

 None

Returns:

Type Description
Dict[str, ndarray]

Dict[str, np.ndarray]: benchmark/episode/returns,
Dict[str, ndarray]

benchmark/episode/length, and result.curve.diagnostics'
Dict[str, ndarray]

own keys, unchanged.

Checks whether url is a full URL ending in a commit SHA.

Parameters:

Name Type Description Default
url str

The URL to validate.

 required

Returns:

Name Type Description
bool bool

True if url has an http(s) scheme and its last path
bool

segment is a 7-40 character lowercase hex SHA.