Skip to content

Agent

Agent: the common interface Experiment trains, and HParams: the base its per-algorithm hyperparameter structs extend.

The concrete agents - PPO, PQN, Dreamer - each subclass Agent, implement train, and carry their own HParams subclass. This module also holds the shared logging helpers (masked_mean, derive_episodic_metrics) and the agent/train/* metric contract every agent's train must return.

agent/train/* is Agent.train's own guaranteed-floor namespace - see its docstring for exactly what's verified common across every navix agent and why (as opposed to agent/diagnostics/*, algorithm- specific and never guaranteed).

Bases: PyTreeNode

Two strategies exist for looking at a run's results, and they trade off against each other:

  • Experiment.run(log_to_wandb=True) (the default) streams metrics to Weights & Biases as training progresses, via log_to_wandb/ log_to_wandb_on_train_end below. This is the slow path - real network I/O, roughly linear in the number of seeds - but gives you wandb's hosted dashboards, run comparison, etc.
  • Experiment.run(log_to_wandb=False) skips wandb entirely and just returns logs (the same pytree these methods consume) directly - no network calls, so it's dramatically faster. Pair it with navix.benchmarks.plotting to get a local matplotlib dashboard from logs instead of a wandb one. See issue #60.

On the base class since every agent needs one - unlike algorithm-specific internals (e.g. PPO's sgd_step).

Deprecated: use log_to_wandb instead.

Deprecated: use log_to_wandb_on_train_end instead.

Streams one training update's metrics to Weights & Biases, deriving the agent/episode/* aggregates from the raw agent/train/* buffers first. A no-op unless logs["agent/train/updates"] is a multiple of hparams.log_frequency. Experiment.run calls this; you rarely call it directly.

Parameters:

Name Type Description Default
logs dict

one update's metrics (a single-update slice of Agent.train's output).

 required
inspectable

optional extra payload for a debug callback.

 None
run

an explicit wandb.Run to log to; defaults to the module-level current run.

 None

Replays log_to_wandb for every recorded update after training has finished - for when train ran fully inside jax.jit and streaming live was not possible. logs here has a leading update axis (the whole history); each kept update is logged in order.

Parameters:

Name Type Description Default
logs dict

the full Agent.train output.

 required
run

an explicit wandb.Run; defaults to the current run.

 None

Trains this agent from scratch and returns logs: the training history every downstream consumer (log_to_wandb, Experiment, navix.benchmarks) reads.

Every navix namespace that ends up in a logs-shaped dict is prefixed by who produced it, not just what kind of thing it is - agent/* here, experiment/* for what Experiment adds after train() already returned, benchmark/* for what Benchmark.summary/details/diagnostics.npz add on top of that (see those modules' own docstrings). Within agent/*, logs' keys split into exactly two namespaces - what's structurally guaranteed (agent/train/*, from the shared collect/derive path every concrete agent already goes through, not something each implementation writes by hand) and what's genuinely bespoke (agent/diagnostics/*) - verified directly against every navix agent's own update/train, not assumed:

  • agent/train/*, guaranteed: agent/train/done_mask/ agent/train/returns/agent/train/lengths (the raw, per-step interaction stream - required by derive_episodic_metrics, which raises KeyError if any is missing) and agent/train/frames/agent/train/updates (identical across every navix agent). None of these are themselves per-episode values - returns/lengths are dense running sums reset on episode boundary, only meaningful where done_mask is true (see navix. environments.environment.Environment.step's info["return"] accumulation).
  • agent/diagnostics/*, bespoke: everything else, including things that look like they should be structural but aren't actually uniform - e.g. an epoch count or learning-rate schedule state exists for PPO/PQN but not in the same shape for Dreamer (three optimizers, not one) - alongside the obviously algorithm-specific values (PPO's agent/diagnostics/entropy/agent/diagnostics/value_loss/ ...; PQN's agent/diagnostics/q_loss/agent/diagnostics/ epsilon; Dreamer's agent/diagnostics/model/*/agent/ diagnostics/actor/*/agent/diagnostics/critic/*). One shared prefix, no shared key names - a caller (e.g. benchmarks/*/*/run.py's TrainingCurve.diagnostics construction) can always filter on key.startswith("agent/diagnostics/") without needing to know which specific keys a given agent happens to log.

Neither of train()'s own two namespaces has anything called agent/episode/* or experiment/costs/*. agent/episode/* is derive_episodic_metrics's own output, computed downstream from agent/train/*'s raw stream. experiment/costs/wall_time/ experiment/costs/fps are added by Experiment.run/ run_hparam_search after train() already returned - real wall-clock timing can't be measured from inside the jax.jit trace train() runs in. A bare agent.train(rng) call, not wrapped by Experiment, returns logs with neither.

Parameters:

Name Type Description Default
rng Array

PRNG key for the whole training run.

 required

Returns:

Type Description
TrainState

Tuple[TrainState, Dict[str, Array]]: The final train
Dict[str, Array]

state, and logs as described above.

Bases: PyTreeNode

Base for every agent's hyperparameter struct (PPOHparams, PQNHparams, DreamerHparams). Holds only the fields common to all; each subclass adds its own (learning rate, rollout length, ...). Frozen - use .replace(...) for a modified copy (this is what the hyperparameter search does).

If True, agents run extra jax.debug callbacks and per-step wandb logging. Slow; off by default.

Log to wandb every log_frequency training updates (1 = every update).

If True, agents also emit an rgb rollout video under render/human in their logs.

Reduces the raw per-step buffers (agent/train/done_mask, agent/train/lengths, agent/train/returns) that Agent.train returns into agent/episode/length, agent/episode/returns, agent/episode/success_rate - one point per training update, masked-mean over completed episodes only. Not itself a per-episode log (nothing in logs is - see Agent.train's docstring): a single aggregate statistic over however many episodes happened to complete that update.

Agent.log_to_wandb computes the same values, but one training update at a time (for live wandb logging); this is the batched equivalent, for reducing an entire already-finished logs history in one call - used by Experiment.run_hparam_search (as the ES search's fitness signal) and navix.benchmarks.plotting (as plot_metric/plot_dashboard's agent/episode/* inputs).

Parameters:

Name Type Description Default
logs Dict[str, Array]

The logs pytree returned by Experiment.run(), Experiment.run_hparam_search(), or a bare Agent.train() call. Must contain agent/train/done_mask/agent/train/lengths/ agent/train/returns, shaped (..., num_steps, num_envs) - any number of leading batch dimensions (e.g. seeds, hparam sets) is supported.

 required

Returns:

Type Description
Dict[str, Array]

Dict[str, Array]: logs, plus agent/episode/length,
Dict[str, Array]

agent/episode/returns and agent/episode/success_rate
Dict[str, Array]

(shape: logs' leading batch dimensions, with num_steps and
Dict[str, Array]

num_envs reduced away). logs itself is not mutated.

Raises:

Type Description
KeyError

If logs is missing agent/train/done_mask, agent/train/lengths, or agent/train/returns.

Mean of values over entries where mask is True, computed via a masked sum/count rather than boolean-indexing values[mask]. Boolean-indexing produces a dynamically-shaped result (its size depends on how many entries are True, which usually varies across calls), forcing JAX to recompile a fresh XLA program for every distinct count it hasn't seen before. This keeps the output shape fixed - values.shape reduced over axis - regardless of how many entries are masked, so XLA compiles it once and reuses it.

Parameters:

Name Type Description Default
values Array

The values to average.

 required
mask Array

A boolean array, broadcastable to values.shape, selecting which entries to include.

 required
axis

The axis or axes to reduce over. None reduces to a scalar.

 None

Returns:

Name Type Description
Array Array

The mean of values where mask is True, reduced over axis.