Logging, evaluation and checkpoints¶

See Logging and Checkpoints and resuming.

ajax.logging.wandb_logging ¶

Helpers for Weights & Biases and TensorBoard logging

LoggingConfig ¶

Pass along the wandb config cleanly

backend `property` ¶

backend

Whether a run writes anywhere (W&B or TensorBoard), so needs the logging worker; without one it still evaluates (train returns the evaluations).

new_run_id ¶

new_run_id()

A fresh run id in W&B's format (wandb.sdk.lib.runid.generate_id: 8 random base-36 digits), without importing wandb.

start_async_logging ¶

start_async_logging()

Spawn the logging process (once; no-op if already running).

stop_async_logging ¶

stop_async_logging()

Drain the queue, then stop the logging process.

init_logging ¶

init_logging(run_id, logging_config, run_seed=None)

Register a run: create it in the main process (so it appears in the W&B UI immediately), then send metadata to the worker so it can flush batches later.

When run_seed is given (one call per seed), it is recorded alongside the config so every run is reproducible and can be paired across conditions by seed (see the per-run config.json written for TensorBoard runs).

vmap_log ¶

vmap_log(log_metrics, index, run_ids)

Forward per-seed metrics to the logging process.

iter_tfrecord ¶

iter_tfrecord(path)

Iterate raw record data from a TFRecord-style TensorBoard event file. (Each record = [length][crc_len][data][crc_data])

load_scalars_from_tfevents ¶

load_scalars_from_tfevents(log_dir)

Reads all scalar summaries from TensorBoard event files using only tensorboardX (no TensorFlow, no tensorboard).

ajax.evaluate ¶

ValueEstimate `dataclass` ¶

ValueEstimate(value, gamma, reward_scale=1.0)

An agent's critic, read for the value bias :func:evaluate returns: value(agent_state, obs, action) is its estimate (one per env) of the return from obs taking action, of the rewards times reward_scale discounted by gamma, as its targets are.

normalised_reward ¶

normalised_reward(x, norm_info)

x, a reward or a sum of rewards in the env's units, in the agent's: divided by the scale the training env's reward normaliser holds (norm_info.reward, frozen), as collection divides each reward; unchanged without one.

setup_environment ¶

setup_environment(
    env,
    env_params,
    num_episodes,
    norm_info,
    gamma,
    action_repeat=1,
    episode_length=None,
)

The training env's stack, rebuilt for evaluation (gymnax or brax).

Ajax's own layers come off (:func:~ajax.environments.create.strip_ajax_wrappers) and go back on as training composed them (:func:~ajax.environments.create.add_ajax_wrappers): the clip to a continuous action space's bounds, the observation normaliser frozen at norm_info (None: no normaliser) and applied to the observations only where training applied it, rewards never normalised. The task under them -- the user's wrappers, the flattening -- is kept as trained on (gymnax).

brax / playground tasks are rebuilt from their id with num_episodes parallel envs. episode_length (simulator steps) and action_repeat are the training env's; the defaults are None -> the env's native episode length, falling back to 1000, and repeat 1. gymnax envs do not support action_repeat > 1.

get_deterministic_action_and_entropy_fn ¶

get_deterministic_action_and_entropy_fn(
    actor_state, recurrent, continuous
)

Return a function mapping (obs, done, hidden) → (action, entropy, hidden).

In recurrent mode the hidden state is an explicit input/output so the eval loop can thread it through its carry (the live carry on actor_state belongs to the training envs and has the wrong batch size here). Non-recurrent mode passes hidden through untouched.

step_environment ¶

step_environment(
    mode,
    env,
    env_params,
    recurrent,
    actor_state,
    continuous,
    expert_policy=None,
    train_frac=None,
    eval_action_transform=None,
    agent_state=None,
    pid_gain_policy=False,
    augment_obs_with_expert_action=False,
    augment_obs_with_expert_state=False,
    value_estimate=None,
    command=None,
    norm_info=None,
)

Return a pure function stepping the evaluation (an :class:_EvalCarry).

The expert acts in agent units (:class:ExpertInAgentUnits), as every action here; the executed one is mapped onto the env's bounds (:func:agent_action_to_env), as in collection.

pid_gain_policy: when True, the actor output is interpreted as PID gain modulation: env_action = expert.step_with_gains(state, obs, anchorexp(ln10a)).

augment_obs_with_expert_action: when True, the actor receives [env_obs, expert_action] as input. The expert is queried statefully (matching training) so the policy sees the same augmented obs at eval that it saw during training.

value_estimate: the critic's estimate at the actor's input and the action collection stores, each step: the executed action in agent units, the actor's (the gains') under pid_gain_policy.

command: a command-conditioned actor's (see :func:evaluate): the actor reads command.actor_input(obs, c), and each step's reward, in the agent's units (norm_info, :func:normalised_reward), and end update c.

step_environment_expert ¶

step_environment_expert(
    mode, env, env_params, expert_policy
)

Step function for the expert's own rollout, called as the agent's rollout calls it (:func:step_environment): a stateful expert (one with init_state) as expert_policy(state, obs) -> (action, state), a stateless one as expert_policy(obs) -> action, in agent units and mapped onto the env's bounds. A discrete expert's action index comes as an (n, 1) column, as a replay keeps it; the env takes (n,).

evaluate ¶

evaluate(
    env,
    actor_state,
    num_episodes,
    rng,
    env_params,
    recurrent=False,
    gamma=0.99,
    norm_info=None,
    avg_reward_mode=False,
    expert_policy=None,
    train_frac=None,
    eval_action_transform=None,
    max_eval_steps=None,
    agent_state=None,
    pid_gain_policy=False,
    augment_obs_with_expert_action=False,
    augment_obs_with_expert_state=False,
    eval_expert_policy=None,
    value_estimate=None,
    command=None,
)

Evaluate the actor's deterministic policy on num_episodes episodes; with expert_policy (the training expert, which builds the actor's input) the expert's return too, that of eval_expert_policy when given, on the same episodes (the same reset and key). The experts speak env units: they are read in agent units here (:class:ExpertInAgentUnits; an agent's are already).

command evaluates a command-conditioned actor (UDRL's desired return and horizon): command.initial(agent_state, num_episodes) is each episode's first command, the actor reads command.actor_input(obs, c), and command.update(c, reward, done) follows each step, the reward in the agent's units, normalised as in training: the units of its targets.

Returns the mean return, entropy, average reward (avg_reward_mode), value bias (value_estimate), episode length and expert's return, each NaN without what it needs. The value bias is the mean over the evaluated steps of the critic's estimate (:class:ValueEstimate) minus the discounted return realised from the step, in the critic's units: the rewards normalised as in training (:func:normalised_reward) and scaled. A recurrent critic's memory is not run here: NaN then.

evaluate_policy ¶

evaluate_policy(
    env_args, policy_fn, init_carry_fn, num_episodes, key
)

Mean return and mean length of one episode in each of num_episodes envs.

The evaluation protocol of the agents that own their collection (the world-model agents, DESIGN §5.4), for a stateful policy with the row collector's protocol policy_fn(carry, obs, is_first, key) -> (action, carry, extras) (extras ignored):

the env is rebuilt through :func:setup_environment with num_episodes parallel envs and the training action_repeat and episode length; actions go through :func:ajax.environments.utils.agent_action_to_env, as in the row collector, so train and eval act on the env identically (the env's clip to its bounds is then a no-op);
the envs are reset with key (the same reset-key derivation as :func:evaluate, so both see the same initial states for one key);
the policy starts from init_carry_fn(num_episodes) (a zero carry) with is_first set on the first step only;
every env runs exactly one episode: a scan over the episode length in agent steps (:func:ajax.environments.utils.agent_episode_length) with rewards and lengths masked after each env's first done.

Preconditions (they raise): the training env does not normalise observations or rewards (the rebuild is the raw env, see :func:ajax.environments.row_collector.check_unnormalized_env), and env_params describes one system (per-env, batched params cannot be rebuilt with num_episodes envs; evaluate a nominal system instead).

ajax.checkpoint ¶

Checkpoint save/restore for Ajax agents.

Serializes the pytree data of a BaseAgentState (arrays, counters, and other leaves) while discarding non-data fields such as optax GradientTransformation closures and flax apply_fn callables. Those closures are rebuilt at load time by re-initializing a skeleton state from the same agent config; the loaded data is then grafted onto the skeleton via flax.serialization.from_state_dict.

Typical resume flow (resume in tests/probing/runs.py runs it): skeleton = agent.train(seed=seeds, n_timesteps=0) # init only if checkpoint_exists(path): skeleton = restore_into(skeleton, path) final = agent.train(..., initial_state=skeleton, ...) save_checkpoint(final, path)

save_checkpoint ¶

save_checkpoint(state, path)

Save the pytree data of state to path. Atomic write.

restore_into ¶

restore_into(skeleton, path)

Load path and graft its data onto skeleton.

skeleton is a freshly-initialized agent state with the same structure as the saved one. The optimizer transformations and apply functions come from skeleton; only the pytree data (params, optimizer mu/nu/count, buffer arrays, rngs, counters) is replaced.