Extensions¶
An extension adds a research feature to any agent whose loop folds the phases it uses. The user guide shows how to use them and Writing an extension how to add one.
ajax.extensions.base ¶
The Extension framework — composable research features for any agent.
An :class:Extension is a research feature (expert guidance, observation
augmentation, online behavioural cloning, target modifiers, exploration
schemes, instrumentation, …) that attaches to any agent without
bloating the agent's own file. It implements only the lifecycle phases
it touches; an agent's training loop folds every extension through each
phase, in list order.
Design contract¶
- Static. Extensions are plain Python objects passed as static
arguments to the jitted training step, so the fold unrolls at trace
time (a Python
forover a fixed tuple — no dynamic dispatch). Every extension must therefore be hashable: the default id-based hash is fine for a one-off instance; a@dataclass(frozen=True)is preferred for instances carrying config, so the JIT cache is reused across equivalent runs. - State lives in the pytree. An extension never stores mutable state
on
self— only frozen config. Anything that changes during training is a pytree returned by :meth:Extension.init_stateand carried inBaseAgentState.ext_state(one entry per extension; the default()means stateless, zero overhead). - No-op by default. Every phase method has a default that is the
identity /
0.0/{}/ unchanged-state, so an extension is exactly as invasive as the phases it overrides, and an empty :class:ExtensionStackis genuinely zero-cost.
Phases¶
init_state build the extension's pytree state (once, on fresh init)
pretrain one-shot, before the training loop (MC/BC pre-training)
on_target transform the TD / value target
critic_loss extra additive critic-loss term (summed over extensions)
actor_loss extra additive actor-loss term (summed over extensions)
action override the collection-time action (None = defer;
the last non-None proposal in the stack wins)
eval_action override the evaluation-time action (same rule)
transform_action make the executed action a function of the policy's
and the expert's (residual RL), in collection, the
losses and evaluation
post_update hook after the update step (φ-refresh, schedules, …)
eval_metrics extra metrics, merged into the eval log
on_target, critic_loss and actor_loss may return (value,
metrics) instead of their value: metrics, a dict of scalars named
<extension>/<metric>, is logged with the update's own metrics (the last
update's, or the mean over the update's epochs and minibatches, as the
agent logs its own). Two extensions logging one name raise.
Phase batches¶
The batch dict of a phase holds the keys :data:PHASE_BATCH_KEYS
lists on every agent that folds it; an agent may add its own (read its
fold call), and an extension raises on a key it needs and does not find,
rather than return a no-op. gamma is None for an average-reward
agent, dones is where the target cut the bootstrap and targets the
target, held constant; the critic (DQN's and PQN's q_state) and the
actor come at the differentiated parameters. The replay agents'
raw_observations is None in a run without an expert, whose replay
does not keep them. a_expert is the agent's expert_policy's action
at each row, stored at collection from the expert's state along the
trajectory; None without one.
ExtensionContext ¶
Shared context handed to every extension phase.
step and rng are traced; total_steps is static so it can
drive Python-level decisions (e.g. schedule lengths) at trace time.
The pretrain phase of a run on ajax.agents.loop.TrainLoop also
gets its static env_args and, off-policy, its update(agent_state,
experience) -> (agent_state, aux) and replay buffer, so an
extension can run the agent's own update; each is None where the run
has none.
Extension ¶
Base class for a composable research feature.
Subclass and override only the phases the feature touches. Keep
subclasses hashable (@dataclass(frozen=True) is the recommended
form) and stateless on self — see the module docstring.
init_state ¶
Return this extension's initial pytree state.
Called once, on fresh initialisation only (skipped on resume).
The default () marks a stateless extension — no pytree
overhead in BaseAgentState.ext_state.
pretrain ¶
One-shot step before the training loop (e.g. MC/BC pre-train).
Returns the (possibly updated) (agent_state, ext_state).
on_target ¶
Transform the TD / value target; or return (target, metrics)
(the module docstring).
critic_loss ¶
Extra additive critic-loss term (summed across extensions); or
(term, metrics).
actor_loss ¶
Extra additive actor-loss term (summed across extensions); or
(term, metrics).
action ¶
Override the collection-time action; None defers.
eval_action ¶
Override the evaluation-time action; None defers.
transform_action ¶
The executed action from the policy's actions and the
expert's a_expert; the policy's by default.
reward ¶
Replace a rollout's rewards before the agent's targets read
them. Returns (rewards, ext_state, metrics).
post_update ¶
Hook after the update step. Returns (agent_state, ext_state).
eval_metrics ¶
Extra metrics, merged into the eval log.
experts ¶
The experts this extension calls itself, on the observations it is given (without a state); none by default.
implemented_phases ¶
The phases this extension actually overrides (vs the no-op base).
bind_to_agent ¶
Return self, optionally populating closed-over agent context.
The default is the identity: an extension ignores every kwarg it
does not need. One that closes over values the agent factory
resolves (env_args, buffer, …; e.g. :class:MCPretrain,
:class:PhiRefresh) overrides this to return a new frozen instance
with them filled in, so a whole stack binds without per-extension
type checks (:meth:ExtensionStack.bind_to_agent).
ExtensionStack ¶
A folded, immutable collection of extensions.
Holds the static tuple of :class:Extension objects and folds them
through each phase in list order. An empty stack is a true no-op.
Hashable (so it can be a JIT static argument) iff its extensions are.
implemented_phases ¶
The phases at least one extension of the stack implements.
fold_init ¶
A fresh run's extension setup: every extension's initial state,
then its one-shot pretrain at step 0 (rng split between the
two; context: :meth:fold_pretrain's). Skipped on resume, where
the state already carries both.
fold_init_states ¶
agent_state with every extension's initial state (one key
each) in its ext_state.
fold_pretrain ¶
One-shot pretrain fold; returns the updated agent_state.
context: the run's static :class:ExtensionContext fields
(env_args, update, buffer).
fold_post_update ¶
Post-update fold; returns the updated agent_state.
fold_reward ¶
Reward fold, in list order, each seeing the previous rewards; returns the state (the extensions' states updated), the rewards and the metrics.
fold_on_target ¶
Target-modifier fold; returns the (possibly reshaped) target and the extensions' metrics.
fold_critic_loss ¶
Critic-loss fold; returns the summed additive critic-loss term and the extensions' metrics.
fold_actor_loss ¶
Actor-loss fold; returns the summed additive actor-loss term and the extensions' metrics.
fold_eval_metrics ¶
Eval-metrics fold; returns the extensions' merged metrics (each
extension gets its own key; the context carries rng).
bind_to_agent ¶
Return a new stack with every extension bound to agent_context
(each :meth:Extension.bind_to_agent picks the kwargs it needs);
each agent factory passes the context it owns. An empty stack
returns self.
with_experts ¶
The stack with convert applied to every expert its extensions
hold (their expert_fields; one still None is left). A
convert memoised by identity keeps one expert one object, which
:func:ajax.extensions.expert.reads_stored_expert_actions compares.
check_extensions ¶
check_extensions(
agent_name,
stack,
capabilities,
phases,
*,
context=(),
pretrain_context=(),
buffer_fields=(),
agent_expert=None,
takes_expert_policy=False,
stores_expert_actions=False,
)
Reject the extensions of stack that cannot run on agent_name,
whose run has capabilities and whose code folds phases, binds the
context keys, fills the pretrain_context fields and stores the
buffer_fields
(:func:~ajax.extensions.capabilities.check_capabilities), or that call
a stateful expert without its state (:func:check_extension_experts).
Raises a ValueError naming every offending extension: one whose phase
an agent never folds would otherwise be silently ignored. Agents call it
at construction (ActorCritic.__init__).
holds_expert ¶
Whether a run has an expert: the agent's expert_policy, or one an
extension holds (its expert_fields). Replay keeps the raw
observations and the expert flags, which only experts read, for such
a run (init_buffer's include_raw_obs, include_is_expert).
check_extension_experts ¶
Reject an extension calling a stateful expert (one with
init_state, called expert(state, obs)) without its state.
An extension calls the experts it declares (:meth:Extension.experts)
on the observations it is given, where no state of the expert exists:
a stateful one is allowed only as the agent's expert_policy of an
agent that stores that expert's actions from its state
(stores_expert_actions), for an extension that reads them
(reads_agent_expert_actions). Stateless experts are always allowed.
phase_output ¶
A training phase's value and metrics: it returns its value, or
(value, metrics) (the module docstring).
merge_metrics ¶
The extensions' training metrics in one dict; a name two of them log raises (two extensions of one class report the same names).
ajax.extensions.expert ¶
Expert-guidance research features as composable :class:Extensions.
- :class:
OnlineBC— value-weighted decaying BC loss term. - :class:
ImitationLoss— plain BC loss term towards the expert. - :class:
ResidualPolicy— executeclip(a_expert + scale·a_pi). - :class:
JSRLCurriculum— per-episode expert→learner handoff. - :class:
ExpertValueMetrics— the critic at the expert's actions, logged.
Each extension carries only frozen config on self; the maths are the
functions below. The expert is SAC's expert_policy: OnlineBC,
ResidualPolicy and JSRLCurriculum read the expert action SAC computes.
ImitationLoss holds its own expert, and reads the actions the agent stored
when that expert is the agent's.
OnlineBC
dataclass
¶
Bases: Extension
Online, value-weighted, warmup-decaying behavioural-cloning term.
A bc_coef · w(s) · ||μ_θ(s) - a_expert||² penalty
(:func:compute_online_bc_loss) that decays to zero once the training
fraction exceeds critic_warmup_frac. It reads the actor's mean
action (pi_mean), the expert's action (a_expert, which SAC puts
in its actor-loss batch) and the observations from the batch, and the
frozen expert critic phi with its value range
(expert_critic_params, expert_v_min, expert_v_max) and the
critic phi is applied with from the agent's state, which
:class:MCPretrain fits (refused without it). The expert's action is
the agent's expert_policy's, so an agent without one refuses it
(requires agent_expert); it raises on a batch without a_expert (an
agent other than SAC). Until phi* exists the term is 0.
ImitationLoss
dataclass
¶
ImitationLoss(
expert_policy,
coef=1.0,
q_filter=False,
anneal_frac=None,
target="expert",
alpha=None,
reads_stored_expert=False,
buffer=None,
action_scale=None,
augment_obs_with_expert_action=False,
name="imitation_loss",
)
Bases: Extension
Behavioural-cloning regulariser: coef · mean((μ(s) - a_E(s))²) / 4.
Same form as the actor-side BC term of TD3+BC (Fujimoto & Gu, 2021),
with a frozen expert's action as the target: μ(s) is the actor's mean action
(pi.mean(), which is TD3's deterministic action), a_E(s) the
expert's action, held constant: its call on the raw observation, or,
when it is the agent's expert_policy, the action collection stored
from its tracked state (which a stateful expert, a PID with its
integral, needs). The mean runs over
batch and action dims; / 4 maps the squared distance of two
actions in [-1, 1] to [0, 1]. Unlike :class:OnlineBC the term
is unweighted and needs no pre-trained expert critic.
q_filter=True counts only the rows where the critic rates the
expert's action above the policy's (:func:q_filter_weight, at
μ(s)) and still divides by every row, as the paper's fixed λ2 =
1/N_D does: the Q-filtered BC of Nair, McGrew, Andrychowicz, Zaremba
and Abbeel (ICRA 2018, Sec. IV-C, eq. 8). Deviations: replay rows
labelled with the expert's action, not a demonstration buffer (the
replay's expert flags would restore it); coef on the BC term, not
the paper's λ1 = 1e-3 on the RL term and λ2; the policy's mean
action and the minimum over the agent's critics (TD3's actor reads its
first, REDQ's the mean of its ensemble), not DDPG's one critic. The
filter needs a Q critic, which SAC, TD3, REDQ, ASAC and AVG put in
their actor-loss batch (critic_state); a state-value critic (PPO,
APO) cannot rate an action. It refuses :class:ResidualPolicy, whose
critic reads the executed action. imitation/q_filter_frac logs, at
each evaluation, the share of a replay batch the filter keeps (every
agent but AVG, which keeps no replay).
With anneal_frac the term is multiplied by a weight falling
linearly from 1 to 0 at that fraction of training (:func:anneal_weight),
as in kickstarting (Schmitt et al., 2018, Kickstarting Deep
Reinforcement Learning, arXiv:1803.03835, eqs. 2-3), whose distillation
weight on the teacher decays so the student ends on its own RL loss.
Deviations from kickstarting: a squared error on the mean action
instead of the cross-entropy between the teacher's and the student's
policies (the expert is deterministic: it gives an action, not a
distribution); a linear schedule set by anneal_frac instead of the
weights population-based training tunes. None keeps the weight at 1.
target="replay" clones the replayed rows' own actions (actions,
which SAC, TD3, REDQ and ASAC put in their actor-loss batch) instead of
an expert's: logged data needs no expert, and expert_policy may be
None. alpha makes the actor loss TD3+BC's (Fujimoto and Gu, 2021),
λ · (the agent's own) + coef · mean((μ(s) - a)²) with λ = alpha /
mean|Q_1(s, μ(s))| (:func:td3bc_lambda), the actions compared in
the env's units and without the / 4: the term adds (λ - 1) times
the agent's own loss (agent_loss in its actor batch), already
counted. On TD3, whose own loss is -mean Q_1, that is the paper's
-λ mean Q + mean((π(s) - a)²); anneal_frac and the filter act
on the cloning part only. With both, on TD3, a transition
dataset trained offline (train_offline) and
normalize_observations=True (the dataset's statistics, sqrt(var +
1e-8) in place of the reference's std + 1e-3), it is TD3+BC.
TD3, REDQ, ASAC, APO, SAC and PPO put pi_mean and
raw_observations in their actor-loss batch, and all but PPO (which
takes no expert) the stored a_expert; an agent whose batch has no
pi_mean (a discrete-action policy) raises at trace time.
reads ¶
The filter and alpha also read the critic from the actor-loss
batch, the replayed target its rows' actions.
bind_to_agent ¶
Read the expert's stored actions when it is the agent's; with
alpha, the env's action half-ranges (TD3+BC compares env actions);
with the filter, bind the agent's replay buffer (its logged share
samples it) and SAC's observation augmentation.
eval_metrics ¶
imitation/q_filter_frac, the share of a replay batch the
Q-filter keeps (an agent with a replay, which it binds).
ResidualPolicy
dataclass
¶
Bases: Extension
Residual RL (Johannink et al. 2019).
The executed action is clip(a_expert + scale·a_pi, -1, 1): SAC's
action pipeline applies :meth:transform_action at collection (its
warmup samples uniformly), and the actor's policy gradient, the
TD-target bootstrap and evaluation flow through the same transform.
Its presence alone turns residual RL on, with SAC's expert policy;
transform_action is a phase only SAC runs, so the other agents
refuse it.
transform_action ¶
clip(a_expert + scale·a_pi, -1, 1) — the residual mix.
build_policy_transform ¶
The actor-loss / TD-target residual transform, (actions,
raw_obs, a_expert_precomputed) -> actions; None (no residual
transform) without an expert_policy.
build_eval_transform ¶
The eval-time residual transform, (raw, expert, obs,
agent_state) -> actions, as :func:ajax.evaluate.step_environment
calls it.
JSRLCurriculum
dataclass
¶
Bases: Extension
True JSRL per-episode handoff (Uchendu et al. 2023).
While step_in_episode < H_t the expert acts; otherwise the
learner does. H_t decays linearly from episode_length to 0
over the first decay_frac fraction of training.
The :meth:action phase owns the per-episode substitution. It reads
post_warmup_action and expert_action off the SAC action
pipeline's batch dict (the same dict :class:EDGEExploration and
:class:~ajax.extensions.target_mods.ValueBox use) and per-env
step_in_episode off agent_state.collector_state; the SAC
pipeline still owns the step_in_episode increment/reset (in
:func:collect_experience). It writes _jsrl_use_expert back, which
the pipeline reads to flag the steps the expert took.
Evaluations log the guide horizon H_t (jsrl/horizon) and the
share of envs the expert drove on the last collection step
(jsrl/expert_frac).
ExpertValueMetrics
dataclass
¶
ExpertValueMetrics(
buffer=None,
expert_policy=None,
augment_obs_with_expert_action=False,
name="expert_value_metrics",
)
Bases: Extension
The critic at the expert's actions, logged at each evaluation, on a
replay batch: value/q_expert_mean, the mean of min_k Q_k(s,
a_E); value/q_gap, that minus the same at the batch's actions
(above 0: the critic rates the expert above what was done); and, once
:class:~ajax.extensions.pretrain.MCPretrain has fitted φ*,
value/phi_star_q_gap_ood, the mean |min_k Q_φ*(s, a_E) - min_k
Q(s, a_E)|.
One replay batch per evaluation pays for them, not a critic pass (two
with φ*) per update for three logged numbers. a_E is the
expert's action stored at collection. A replay agent with an
expert_policy (SAC, TD3, REDQ, ASAC) binds its replay buffer, its
expert and its observation augmentation.
residual_action_transform ¶
Residual RL: executed action is clip(a_expert + scale * a_pi, -1, 1).
The critic was trained on (s, executed_action) tuples, so the policy
gradient must flow through Q(s, executed_action), not Q(s, a_pi).
scale follows Johannink et al. 2019; values < 1 keep the initial
behaviour close to the expert when the policy is randomly initialised.
online_bc_weight ¶
The online BC term's schedule: 1 at the start, falling linearly to 0
at critic_warmup_frac of training.
compute_online_bc_loss ¶
compute_online_bc_loss(
pi_mean,
a_expert,
critic_state,
expert_critic_params,
observations,
train_frac,
critic_warmup_frac,
expert_v_min,
expert_v_max,
bc_coef,
)
Online decaying BC term — value-weighted, warmup-decaying.
Decays to zero after train_frac exceeds critic_warmup_frac. Value-weighting ensures BC is strongest in high-value states.
reads_stored_expert_actions ¶
Whether an extension that carries an expert of its own (IBRL,
LCBGatedBootstrap, CriticBlend, PhiRefresh) reads that expert's actions
as collection stored them (next_a_expert): when it is the agent's
expert_policy, whose state collection tracks along each trajectory.
Another expert is called on the raw observations the extension is
given, without a state (a stateful one is refused at construction,
check_extension_experts).
expert_next_actions ¶
The extension's expert's actions at the replayed rows' s': the
stored ones (reads_stored_expert_actions), else its call on the
raw next observations.
anneal_weight ¶
1 at the start of training, falling linearly to 0 at anneal_frac
of total_steps, 0 after.
q_filter_weight ¶
The Q-filter of Nair et al. 2018 (eq. 8), (batch, 1): 1 where the
critic's ensemble minimum rates the expert's action above the policy's,
min_k Q_k(s, a_E) > min_k Q_k(s, π(s)), else 0; held constant.
td3bc_lambda ¶
TD3+BC's lambda = alpha / mean|Q_1(s, pi(s))|, held constant,
with Q_1 the first critic, the one TD3's actor follows
(sfujim/TD3_BC@8791ad7, TD3_BC.py:141-143: lmbda = alpha /
Q.abs().mean().detach(), actor loss -lmbda * Q.mean() +
mse_loss(pi, action)).
ajax.extensions.pretrain ¶
One-shot pre-training research features as composable :class:Extensions.
CloningPretrain, MCPretrain and BellmanPretrain run once, before the
training loop (the pretrain phase, skipped on resume); PhiRefresh
refreshes MCPretrain's critic during training.
- :class:
CloningPretrain— behaviour cloning of the actor on an expert's rollouts. - :class:
MCPretrain— regress a frozen expert critic on Monte-Carlo returns, then (optionally) lightly nudge the online critic toward it. - :class:
BellmanPretrain— Bellman-bootstrapped critic pre-training on the expert buffer. - :class:
PhiRefresh— periodic self-consistent refresh of the frozen expert critic during training (thepost_updatephase). - :class:
PriorData— a transition dataset in the replay, and the agent's own updates on it before acting.
MCPretrainAux ¶
Diagnostics logged after MC critic pretraining.
PretrainLosses ¶
A pretraining fit's first and last loss, kept in its extension's state
for the evaluations' log. NaN before the fit, when it takes no step, and
in a run resumed from a checkpoint saved before they were kept
(added_fields, :func:ajax.checkpoint.restore_into).
MCPretrainResults ¶
Bases: PretrainLosses
MCPretrain's fit: its first and last regression loss, and the mean
over the fit's rows of φ*'s value at the expert's action (the
smallest of its critics').
CloningPretrain
dataclass
¶
CloningPretrain(
expert_policy,
n_steps=100000,
epochs=10,
batch_size=64,
skip_actor=False,
reset_log_std=False,
reset_actor_head=False,
env_args=None,
actor_optimizer_args=None,
augment_obs_with_expert_action=False,
augment_obs_with_expert_state=False,
use_train_frac=False,
name="cloning_pretrain",
)
Bases: Extension
Behaviour cloning of the actor on an expert's rollouts, before training.
Rolls expert_policy out for n_steps steps per env, observing as
the agent does, then fits the actor to its actions on observations
standardised with the dataset's statistics
(:func:ajax.agents.cloning.pre_train: epochs passes in batches of
batch_size, at the actor's learning rate). The agent's observation
statistics start from the dataset's. skip_actor leaves the actor as
initialised; after cloning, reset_log_std puts its log-std (head
or parameter) back at its initial values and reset_actor_head also
its output head (the mean layer, or a PID actor's gains). Evaluations
log the mean loss of cloning's first and last epoch
(cloning/first_loss, cloning/last_loss).
The fields after reset_actor_head are the agent's context, filled by
:meth:bind_to_agent (SAC, REDQ, TD3 and APO do).
bind_to_agent ¶
Copy the agent's env_args, actor_optimizer_args and
observation layout (augment_obs_with_expert_action,
augment_obs_with_expert_state, use_train_frac). Raises
ValueError when the layout carries the agent's expert's action or
state and expert_policy is another expert, or when another
pretrain-phase extension comes before it in the stack.
pretrain ¶
Clone the actor on ctx.rng (:meth:clone) and keep its losses.
clone ¶
agent_state with its actor cloned on rollouts drawn with
key and its observation statistics seeded, and the mean loss of
cloning's first and last epoch (NaN with skip_actor); raises
ValueError when the agent did not bind its context.
MCPretrain
dataclass
¶
MCPretrain(
expert_policy,
n_mc_steps=10000,
n_mc_episodes=100,
n_steps=5000,
use_online_light=True,
online_light_steps=500,
online_light_lr_scale=0.1,
env_args=None,
network_args=None,
critic_optimizer_args=None,
num_critics=2,
gamma=None,
reward_scale=None,
use_train_frac=False,
augment_obs_with_expert_action=False,
expert_state_aug_dim=0,
use_phi_refresh=False,
name="mc_pretrain",
)
Bases: Extension
Monte-Carlo critic pre-training.
Regresses a frozen expert critic φ* on unbiased MC returns from
expert rollouts, persists φ* + the value range v_min/v_max on
the agent state, and — when use_online_light is set — applies a
weak supervised nudge of the online critic toward φ* to reduce
seed-to-seed variance at initialisation.
Building φ* needs the agent's env, critic network and optimiser
config, gamma and reward_scale: the fields after the
hyperparameters, which the agent fills through :meth:bind_to_agent
(only SAC does). φ* trains with that optimiser config as it is, so
the agent's weight_decay (like its gradient clipping and learning-
rate schedule) also decays φ* while it fits the MC returns.
use_phi_refresh keeps φ*'s optimiser state for a
:class:PhiRefresh in the same stack. Evaluations log the value range
(mc_pretrain/v_min, mc_pretrain/v_max) and the fit
(mc_pretrain/first_loss, mc_pretrain/last_loss,
mc_pretrain/q_expert_mean: :class:MCPretrainResults).
bind_to_agent ¶
Copy the agent's context (env_args, network_args,
critic_optimizer_args, num_critics, gamma,
reward_scale, use_train_frac,
augment_obs_with_expert_action, expert_state_aug_dim) onto a
new instance, and keep
φ*'s optimiser when the stack (extensions) holds a
:class:PhiRefresh.
pretrain ¶
Build the frozen expert critic φ* and persist it on agent_state:
a fresh critic regressed on unbiased MC returns sets
expert_critic_params / expert_v_min / expert_v_max,
then (use_online_light) the online critic is nudged toward it.
The fit's results are the extension's state.
Raises ValueError when the agent did not bind its context.
BellmanPretrain
dataclass
¶
BellmanPretrain(
n_steps=5000,
buffer=None,
gamma=None,
reward_scale=None,
recurrent=False,
update_value_fn=None,
update_target_fn=None,
augment_obs_with_expert_action=False,
name="bellman_pretrain",
)
Bases: Extension
Bellman-bootstrapped critic pre-training.
n_steps critic updates on the replay buffer before training, which
then holds the agent's expert prefill. Bootstraps from an untrained
critic, so the pre-trained values are biased (:class:MCPretrain
regresses unbiased returns instead); kept for ablation parity.
Evaluations log the critic loss of its first and last update
(bellman_pretrain/first_loss, bellman_pretrain/last_loss).
The fields after n_steps are the agent's context, filled by
:meth:bind_to_agent (only SAC does).
bind_to_agent ¶
Copy the agent's buffer, gamma, reward_scale, recurrent
and critic updates (update_value_fn, update_target_fn). It
requires agent_expert: without one the buffer holds no expert
prefill to pre-train on.
pretrain ¶
Pre-train the critic and keep its losses; raises ValueError
when unbound.
PhiRefresh
dataclass
¶
PhiRefresh(
expert_policy,
interval=500,
steps=20,
buffer=None,
gamma=None,
reward_scale=None,
augment_obs_with_expert_action=False,
reads_stored_expert=False,
name="phi_refresh",
)
Bases: Extension
Periodic self-consistent refresh of the frozen expert critic.
Every interval steps, run steps self-consistent Bellman
updates on expert-flagged buffer transitions, keeping φ* aligned
with the live data distribution.
Requires :class:MCPretrain (which creates the refreshable φ*;
refused without it).
buffer, gamma and reward_scale come from the agent through
:meth:bind_to_agent (only SAC binds them); unbound, post_update
raises ValueError.
bind_to_agent ¶
Copy the agent's buffer, gamma, reward_scale and
observation augmentation, and whether its expert is the agent's
(:func:reads_stored_expert_actions).
PriorData
dataclass
¶
Bases: Extension
Prior data in the replay, and the agent's own updates on it before acting (offline-to-online RL).
Adds dataset's rows to the agent's replay, where they stay
(:func:ajax.buffers.dataset.fill_replay), and to its observation and
reward normalisers, each row counted once as collection counts its own;
then runs offline_updates of the agent's update on that replay: its
critic, actor and target steps and its extensions' loss terms, but no
post_update between them. offline_updates=0 is RLPD's prior
data in the replay (Ball et al. 2023) without its symmetric sampling.
Ajax's warm-up acts uniformly: pair it with learning_starts=0.
Needs the replay and update an off-policy loop hands the pretrain
phase (pretrain_context): any agent but SAC, TD3, REDQ, ASAC and DQN
refuses it when built.
pretrain_critic_mc ¶
pretrain_critic_mc(
agent_state,
expert_critic_state,
expert_policy,
mode,
env_args,
gamma,
reward_scale,
n_mc_steps=10000,
n_mc_episodes=500,
n_steps=5000,
batch_size=256,
max_timesteps=None,
augment_obs_with_expert_action=False,
augment_obs_with_expert_state=False,
)
Pre-train critic using Monte Carlo returns from expert trajectories.
Unlike Bellman pretraining (bootstraps from an untrained critic → biased), MC returns G_t = Σ γ^k r_{t+k} are unbiased estimates of V^expert(s). The critic starts with accurate Q(s, a_expert) near the target from step 1.
Collection strategy: single call with n_mc_steps * n_mc_episodes // n_envs timesteps. The n_envs parallel environments reset to different initial states on each episode boundary, giving the same state-space coverage as separate per-seed rollouts — without any mapping over traced keys (which fails inside the outer vmap over seeds). Total transitions ≈ n_mc_steps * n_mc_episodes regardless of n_envs.
(agent_state, frozen_expert_params, obs_batched, action_batched,
| Type | Description |
|---|---|
Tuple[Any, FrozenDict, Array, Array, MCPretrainAux, LoadedTrainState]
|
MCPretrainAux, expert_critic_state_trained) |
pretrain_critic_online_light ¶
pretrain_critic_online_light(
agent_state,
obs_batched,
action_batched,
n_steps=500,
lr_scale=0.1,
)
Weak supervised nudge of the online critic toward φ*.
Goal: reduce seed-to-seed variance at initialization. NOT a full MC pretrain — just reduces starting point spread. Target network stays random and untouched.
refresh_phi_star ¶
refresh_phi_star(
agent_state,
buffer,
phi_refresh_steps,
gamma,
reward_scale,
ext,
augment=False,
)
Periodic self-consistent φ refresh using expert-flagged buffer transitions.
augment: φ reads observations augmented with SAC's expert actions;
ext's expert gives the action at s' (expert_next_actions).
Target: r + γ * min_k Q_φ(s′, π(s′)) — φ* supervises its own bootstraps. Non-expert transitions are masked out, so the gradient comes only from (s, a_expert, r, s') rows where EDGE fired the expert action.
pretrain_critic_bellman ¶
pretrain_critic_bellman(
agent_state,
recurrent,
gamma,
reward_scale,
buffer,
n_steps,
update_value_fn,
update_target_fn,
augment=False,
)
n_steps Bellman critic updates (the agent's update_value_fn)
on replay batches, each followed by a target update at tau 5e-4;
augment: the critic reads observations augmented with the expert's
stored actions, as in SAC's update. Returns the agent state and each
update's critic loss.
ajax.extensions.target_mods ¶
TD-target modifier research features as composable :class:Extensions.
These reshape the SAC Bellman target before it enters the critic loss
(the on_target phase):
- :class:
IBRL— add the positive gapγ(1-d)·max(Q_expert - Q_policy, 0)so the value function matches an argmax action-selection policy. - :class:
LCBGatedBootstrap— soft-blend the policy and expert TD targets by an LCB-scored sigmoid gate. - :class:
CriticBlend— warmup-decaying blend of the Bellman target with the frozen expert critic's. - :class:
MCVarianceCorrection— replace high-ensemble-variance Bellman targets with the MC-pretrained oracle. - :class:
ValueBox— value-threshold expert action override during collection (theactionphase).
The batch argument to on_target is the dict loop.critic_step
builds for SAC: observations, actions, next_observations,
rewards (unscaled), reward_scale, dones, gamma,
raw_next_observations (the environment's own next observations, which
an extension's expert reads: SAC keeps them when an extension carries an
expert), next_a_expert, and SAC's own rng_key, q_preds and
recurrent.
IBRL
dataclass
¶
Bases: Extension
IBRL bootstrap (Hu, Mirchandani and Sadigh, Imitation Bootstrapped Reinforcement Learning, 2023).
Adds γ(1-d)·max(min Q_target(s', a_expert') - min Q_target(s', a_π'),
0) to the Bellman target so the value function is consistent with an
argmax(Q_expert, Q_policy) action-selection rule.
bind_to_agent ¶
Read the expert's stored actions when it is the agent's.
LCBGatedBootstrap
dataclass
¶
LCBGatedBootstrap(
expert_policy,
lcb_beta_init=1.0,
lcb_beta_decay_k=2.0,
lcb_temperature=1.0,
reads_stored_expert=False,
name="lcb_gated_bootstrap",
)
Bases: Extension
LCB-gated bootstrap.
Scores the policy and expert next-actions by an LCB rule and
soft-blends the two TD targets with p_expert =
σ((score_e - score_p)/lcb_temperature). β anneals over training;
evaluations log it (lcb_bootstrap/beta).
CriticBlend
dataclass
¶
CriticBlend(
expert_policy,
critic_warmup_frac=0.15,
reads_stored_expert=False,
name="critic_blend",
)
Bases: Extension
Warmup-decaying blend of the Bellman target with the expert's.
(1-α)·y_bellman + α·y_expert with y_expert = r + γ(1-d)·V_E(s'),
the one-step target on the frozen expert critic V_E(s') = φ*(s',
a_E(s')), so the blend stays a Bellman target (V_E(s') alone has
no reward, discount or terminal mask). α decays linearly to 0 over
critic_warmup_frac of training; evaluations log it
(critic_blend/expert_weight). Requires the frozen expert critic
:class:MCPretrain fits (refused without it).
MCVarianceCorrection
dataclass
¶
Bases: Extension
Replace high-variance Bellman targets with the MC oracle.
Where the inter-critic variance exceeds threshold the Bellman
target is swapped for the MC-pretrained expert critic's estimate.
Requires :class:MCPretrain (refused without it).
ValueBox
dataclass
¶
Bases: Extension
Value-threshold expert-action override.
During collection, override the policy action with the expert's
whenever the frozen-expert value V_expert(s) exceeds a curriculum
threshold that ramps from v_min to v_max over training.
Requires :class:MCPretrain (which supplies phi*, v_min and
v_max; refused without it).
The :meth:action phase reads the SAC action pipeline's per-step dict
(env_action, expert_action, the action of SAC's
expert_policy, at which it also measures the expert's value, and
observations, the critics' input, augmented when it is) and the box
bounds off the agent state (so a resumed run keeps them), and writes
in_value_box, entry_bonus and box_entry back into the dict
for the pipeline to record on the transition. An entry ends the agent's
episode: the expert's value, paid as the bonus, replaces the bootstrap,
whatever its sign. The override (:func:box_action_override) comes
after the warmup / post-warmup choice (action_slot).
Evaluations log the threshold (value_box/threshold), the share of
envs inside the box and the share the expert drove on the last
collection step (value_box/in_box_frac, value_box/expert_frac);
the bounds are :class:MCPretrain's (mc_pretrain/v_min,
mc_pretrain/v_max).
action ¶
Override with a_expert inside the value box, its bounds the
MC-pretrain v_min/v_max on :class:SACState; writes
in_value_box, entry_bonus and box_entry into obs for
the pipeline's bookkeeping (reward shaping, the entry's
termination, is_expert_flag).
threshold ¶
The box's threshold at step: MCPretrain's v_min rising
linearly to its v_max over the run.
blend_modify_target ¶
Blended Bellman: (1-alpha)y_bellman + alphay_expert.
mc_correction_modify_target ¶
mc_correction_modify_target(
target_q,
critic_state,
critic_params_mc,
observations,
actions,
q_var,
mc_variance_threshold,
)
Replace high-variance Bellman targets with MC-pretrained oracle estimate.
box_compute_threshold ¶
Curriculum threshold: v_min + (v_max - v_min) * train_frac.
box_entered ¶
Where the step enters the value box: inside it now, not before.
box_compute_state ¶
Compute box membership and entry bonus, the expert's value
V(s) = min Q_phi*(s, a_exp) at its action a_exp.
Returns (in_box, entry_bonus, v_box).
box_action_override ¶
Override with expert action inside the value box.
discounted ¶
The target's discount gamma, which a bootstrap bonus scales;
refused on an average-reward agent (ASAC, APO: gamma None), whose
target has no discount and for which the bonus has no form.
ajax.extensions.adversarial ¶
Adversarial imitation as a composable :class:Extension.
- :class:
AdversarialReward— GAIL's discriminator reward.
Discriminator ¶
Bases: Module
Logits of "an expert's" for an (observation, action) row.
AdversarialReward
dataclass
¶
AdversarialReward(
demonstrations,
n_updates=2,
learning_rate=0.001,
architecture=("32", "relu", "32", "relu"),
name="gail",
)
Bases: Extension
GAIL (Ho and Ermon, NeurIPS 2016): the policy learns from a reward a
discriminator gives, trained to tell its rollouts from the expert's
demonstrations. Following HumanCompatibleAI/imitation@e5ef188
(algorithms/adversarial/gail.py and common.py): each rollout's rewards
become -log(1 - D(s, a)) from the current discriminator (-
logsigmoid(-logits)), then the discriminator takes n_updates
steps of binary cross-entropy, the demonstrations labelled 1, on the
rollout's rows and as many demonstration rows drawn uniformly. The
environment's own reward is not used. The discriminator reads the action
the env executes, clipped to [-1, 1] in agent units
(:func:~ajax.environments.utils.agent_action_to_env), as the reference
records it (its BufferingWrapper sees the action SB3 steps the env
with, clipped to the box). Its 32-32 ReLU network starts as
torch's nn.Linear and steps by torch's Adam (learning rate 1e-3),
the reference's BasicRewardNet and optimiser. Deviations: the
discriminator reads (s, a) unnormalised, without the reference
examples' running normalisation; its generator batch is the whole rollout, the
reference's demo_batch_size rows of its last one; PPO stands in for
the paper's TRPO, as in the reference. Logs gail/reward,
gail/disc_loss and the shares of expert and rollout rows it
classifies right, gail/expert_acc and gail/policy_acc.
ajax.extensions.offline ¶
Offline-RL critic regularisers as composable :class:Extensions.
- :class:
ConservativeQ— CQL(H)'s conservative critic term. - :class:
CalQL— the same term, its policy-action values floored at the data's Monte-Carlo returns (Cal-QL). - :class:
TargetBC— ReBRAC's penalty on the target's next actions.
ConservativeQ
dataclass
¶
Bases: Extension
CQL(H)'s critic term (Kumar, Zhou, Tucker and Levine, NeurIPS 2020,
eq. 4), in the importance-sampled form of the authors' code
(aviralkumar2907/CQL@d67dbe9, d4rl/rlkit/torch/sac/cql.py:234-264,
min_q_version 3)::
alpha * mean( logsumexp_j [Q(s, a_j) - log density(a_j)] - Q(s, a) )
over the critics and the batch, with a the row's action and the
a_j n_actions uniform actions in [-1, 1]^A and n_actions
of the policy at s and at s', all scored at s. It pushes the
critic down on actions away from the data's and up on the data's. The
reference sums each critic's Bellman loss and term; Ajax's critics share
one mean squared error, so the term is the mean over them, at the same
ratio. It needs a policy with a density: a deterministic one (TD3's) is
a point mass, whose samples are copies of one action and whose density
is undefined. Deviations from the reference's SAC: its target keeps SAC's
entropy term (the reference's deterministic_backup drops it). Logs
cql/term and cql/gap, the logsumexp minus the data's value.
floor ¶
The policy actions' values as the term reads them (CalQL floors them); no metrics.
CalQL
dataclass
¶
Bases: ConservativeQ
Cal-QL (Nakamoto, Zhai, Singh, Mark, Ma, Finn, Kumar and Levine,
NeurIPS 2023): :class:ConservativeQ's term with the policy actions'
values floored at the row's Monte-Carlo return to go before the
logsumexp (nakamotoo/Cal-QL@ac6eafe, JaxCQL/conservative_sac.py:227-240),
so the critic is pushed down no further than the data's own return. The
return to go is the transition dataset's (return_to_go), at its
gamma, in the env's reward units: the agent's gamma and a
reward_scale of 1 match it. Deviation: the reference computes a
collected episode's returns when it ends; collected rows here keep
-inf, no floor. Also logs calql/bound_frac, the share of policy
actions at s whose value the floor raised.
TargetBC
dataclass
¶
Bases: Extension
ReBRAC's critic penalty (Tarasov, Kurenkov, Nikulin and Kolesnikov,
NeurIPS 2023; tinkoff-ai/ReBRAC@694385b, src/algorithms/rebrac.py:
134-144): the bootstrap less beta times the squared distance from
the next actions the target scored (TD3's smoothed next_actions) to
the dataset's next action a' (next_action)::
y -= gamma (1 - d) beta sum_i (next_actions_i - a'_i)^2
A row whose a' is unknown (an episode's last row, a collected one:
NaN) keeps its target. Logs target_bc/penalty, the mean penalty, and
target_bc/known_frac, the share of rows with an a'.
policy_samples ¶
n actions of the actor at each observation, (n, B, A), and
their log-densities, (n, B, 1), held constant.
sampled_q ¶
Every critic at each observation and each of its (n, B, A)
actions: (critics, n, B, 1).
ajax.extensions.exploration ¶
Exploration research features as composable :class:Extensions.
EDGE (Expert Decayed Guided Exploration) substitutes the expert action
for the policy action during collection, gated by a value comparison
that decays over training. :class:EDGEExploration selects one of six
gates with gate: "fixed" / "argmax" / "boltzmann"
(value-gap gates) or "lcb" / "argmax_lcb" / "thompson"
(quality-aware gates); the gate maths are the functions below.
The SAC action pipeline pre-computes the per-step quantities every gate
needs (policy_action, expert_action, observations, ...) and
passes them through the obs argument as a dict, as the batch dict
of the on_target / actor_loss phases. The extension picks an
action (or returns None to defer); the pipeline owns the bookkeeping
(is_expert_flag, buffer_action, expert-state threading). Gate
randomness is threaded through obs["gate_rng"] (read and overwritten
in place), the pipeline's own stream, not the framework's per-extension
rng split (stack._keys).
EDGEExploration
dataclass
¶
EDGEExploration(
gate="fixed",
decay_frac=0.3,
tau=1.0,
fixed_prob=0.5,
lcb_beta_init=1.0,
lcb_beta_decay_k=2.0,
lcb_temperature=1.0,
lcb_asymmetric=False,
epsilon_floor=0.0,
name="edge_exploration",
)
Bases: Extension
Expert Decayed Guided Exploration during collection.
gate selects the substitution rule (see module docstring). The
other fields:
decay_frac— value-gap gates decay to 0 over this fraction.tau— Boltzmann temperature on the value gap.fixed_prob— substitution probability for the fixed gate.lcb_beta_init/lcb_beta_decay_k— LCB pessimism schedule.lcb_temperature— softmax temperature of the LCB / Thompson gate.lcb_asymmetric— LCB on the expert arm, UCB on the policy arm.epsilon_floor— minimum policy-pick probability (Thompson).
The expert is SAC's expert_policy: the gate compares the action it
gave the pipeline with the policy's. Evaluations log the share of envs
the expert drove on the last collection step (edge/live_expert_frac)
and the gate's schedule: edge/beta for the LCB gates, edge/decay
for the fixed, argmax and Boltzmann gates (the Thompson gate has none).
needs ¶
The value gates compare the critic's values of the two actions; the LCB and Thompson gates read its ensemble.
action ¶
Apply the configured EDGE gate to substitute the policy action.
obs is the SAC action pipeline's per-step dict (the module
docstring): policy_action, expert_action, observations
(the critics' input, augmented with the expert-action dims under
augment_obs_with_expert_action), critic_params and
critic_state (the frozen expert critic if MC pretraining ran,
else the live critic), post_warmup_action and gate_rng.
Side effects: writes _edge_use_expert back into obs for
the pipeline's is_expert_flag bookkeeping, and updates
obs["gate_rng"] for any downstream gate.
eval_metrics ¶
The fraction of envs an expert drove on the last collection step, and the gate's schedule at the collector's timestep.
lcb_beta ¶
The LCB pessimism init (1 - t / T)^decay_k, annealed over the
run (t / T clipped to [0, 1], T at least 1).
lcb_score ¶
min Q - beta (max Q - min Q) over the ensemble axis: the critics'
disagreement as epistemic uncertainty, penalised.
edge_compute_value_gap ¶
Compute Q(s,pi*) - Q(s,pi). Positive means expert is still better.
Returns (gap, q_policy) — q_policy is reused by the Boltzmann gate.
edge_compute_decay ¶
Linear decay: 1 -> 0 over decay_frac of training.
decay_frac=0.0 means never decay (gate stays active for full training).
edge_argmax_gate ¶
Deterministic: use expert whenever gap > 0 and decay > 0.
Returns (use_expert_mask, unchanged_rng).
edge_boltzmann_gate ¶
Adaptive: p = decay * sigmoid(gap / (tau * |Q|)).
Returns (use_expert_mask, updated_rng).
edge_fixed_gate ¶
Fixed probability: p = decay * fixed_prob.
Returns (use_expert_mask, updated_rng).
edge_lcb_scores ¶
edge_lcb_scores(
obs,
policy_action,
expert_action,
critic_state,
critic_params,
beta,
asymmetric=False,
)
Each candidate's score (score_expert, score_policy): its
:func:lcb_score, the critics' disagreement penalising uncertain
(OOD) candidates; the gate then compares the two.
asymmetric scores the policy arm optimistically instead (UCB,
max Q + beta (max Q - min Q)), so the gate hands control to the
policy where its critic is uncertain, which is where fresh on-policy
Bellman anchors help most; the expert arm keeps its LCB, staying
conservative about following an unreliable expert. Symmetric LCB
over-penalises the policy exactly where it explores beyond the
expert's tube.
edge_lcb_argmax_gate ¶
Deterministic argmax gate over LCB scores.
use_expert = score_e > score_p.
Compared to edge_lcb_gate (softmax), this collapses the gate to
a hard threshold while preserving the LCB scoring (mu - beta*sigma)
so the (gating form, scoring rule) ablation can vary one axis at
a time. Used by the argmax_lcb corner of the 2x2.
edge_lcb_gate ¶
Stochastic Boltzmann gate over LCB scores.
p(use_expert) = sigmoid((score_e - score_p) / temperature).
Hard threshold (argmax) creates a self-fulfilling expert preference: expert always wins → buffer all expert → critic never learns Q for policy actions → expert keeps winning. The stochastic version ensures policy occasionally gets sampled, breaking the inertia.
Temperature → 0 recovers argmax; → ∞ recovers uniform mixing.
edge_compute_thompson_stats ¶
Per-action Q mean / std across critic ensemble: (mu_e, sigma_e, mu_p, sigma_p).
edge_thompson_gate ¶
Thompson-sampling gate, with optional epsilon floor on policy picks.
Treat each candidate's Q as Gaussian(mu, (temperature * sigma)^2), draw one sample per side, pick whichever is larger. Equivalent to p(use_expert) = Phi((mu_e - mu_p) / (T * sqrt(sigma_e^2 + sigma_p^2))) in expectation, but stochastic per-step (which is what we want for exploration: the gate itself injects noise rather than deferring to a fixed rule).
Symmetric: high uncertainty on either side widens the gate, never biasing it. Confident estimates dominate. With both sigmas → 0 we recover deterministic argmax.
epsilon_floor: with this probability, force the policy regardless of the Thompson outcome. Required when the expert dominates by many sigmas (Phi -> 1 -> Thompson never picks the policy -> the critic never sees policy actions in its updates -> Q for policy actions stays stale -> the gate can never flip even if the policy improves). Set carefully on brittle envs: too high crashes the agent, too low starves the policy of evaluation data. ~0.001-0.05 is a sane range.
ajax.extensions.ensemble ¶
Critic-ensemble research features as composable :class:Extensions.
- :class:
KernelRepulsion— SVGD-style function-space repulsion between the critics of an ensemble (thecritic_lossphase).
KernelRepulsion
dataclass
¶
Bases: Extension
coef times :func:q_kernel_repulsion of the critics on the
critic-loss batch, added to the critic loss (REDQ's SVGD variant).
It reads the batch every replay agent's critic step passes
(observations, actions, critic_params, critic_state),
so it costs one more ensemble forward. Feedforward critics only:
agents with memory reject extensions, so repulsion with recurrent
memory is unsupported.
q_kernel_repulsion ¶
The mean RBF kernel between the critics' Q-functions on a batch.
q_preds is (num_critics, batch, ...); each critic's outputs
form one feature vector. The bandwidth follows the median heuristic
(Liu & Wang 2016, SVGD), held constant so the kernel adapts to the
current spread without a confounding gradient; minimising the mean
kernel pushes the critics apart in function space.
ajax.extensions.instrumentation ¶
EVarEst measurement / instrumentation as composable :class:Extensions.
The measurement code of EVAREST's evarest/rl/, foldable into any Ajax
agent (SAC, DQN, REDQ, TD3 off-policy; PPO, PQN, APO, AVG on-policy):
extensions=[ConditioningMetrics(), BiasVoreDecomposition(...), ...].
Coverage¶
- :class:
ConditioningMetrics— eval_metrics phase, plasticity diagnostics (srank, dormant fraction, feat / weight L2 norms). Reads the agent's replay buffer (off-policy) ORagent_state.last_rollout(on-policy withexpose_recent_rollout), auto-detected. - :class:
BiasVoreDecomposition— eval_metrics phase, EVarEst bias^2 / Var(residual) / MSE / |bias| decomposition of the critic objective. Mirrors the math in EVAREST'sloss.py/mechanism.py. - :class:
BiasVorePenalty— critic_loss phase, the EVarEst critic objective itself.penalty(alpha) = ((1-2alpha)/alpha) * Var(q_preds - target_q)added on top of the agent's vanilla MSE.
Without a measurement batch (no replay buffer, no last_rollout yet)
the eval-metric extensions emit no metrics rather than wrong numbers. The
maths are intentionally kept identical to the EVAREST sources.
ConditioningMetrics
dataclass
¶
ConditioningMetrics(
buffer=None,
on_policy_batch=256,
critic_loc="auto",
concat_action="auto",
name="conditioning_metrics",
)
Bases: Extension
Plasticity diagnostics on a state batch sampled from the agent's visited-state distribution.
Reads the critic's encoder features:
* critic_loc='critic' + concat_action=True ⇒ Q(obs,act)
ensemble critic on critic_state.
* critic_loc='critic' + concat_action=False ⇒ V(obs)
ensemble critic on critic_state.
* critic_loc='actor' ⇒ Q-network on actor_state (DQN, PQN),
penultimate features read via the encoder_features sow.
'auto' (the default for both) reads them from the networks, never
from the agent's name or the batch source: the critic branch when
critic_state's network is an ensemble critic (it exposes
apply_encoder_ensemble), the actor branch otherwise; on the critic
branch, the action is concatenated when the critic's parameters were
built on the observation and action (:func:_takes_action).
Pass buffer=agent.buffer for off-policy agents whose buffer object
isn't on the agent state itself; on-policy agents must set
BaseAgentConfig.expose_recent_rollout=True.
BiasVoreDecomposition
dataclass
¶
BiasVoreDecomposition(
buffer=None,
on_policy_batch=256,
gamma=0.99,
reward_scale=1.0,
value_target_fn=None,
name="bias_vore_decomposition",
)
Bases: Extension
EVarEst bias^2 / Var(residual) / MSE / |bias| decomposition.
Ports the math from EVAREST loss.py /
mechanism.make_evarest_decomposition_metrics. The residual is
reconstructed from the agent's current critic on a measurement
batch sampled the same way :class:ConditioningMetrics samples
(replay buffer for off-policy, last_rollout for on-policy).
By default it is the critic's one-step TD residual under the current
policy, prediction - (reward_scale * r + gamma (1 - d) next), with
the critic read from the networks (:func:_critic_kind):
- Q(s, a) ensemble:
next = min_i Q_i(s', a') - alpha log pi(a'|s'),a' ~ pi(s'),alphathe agent's temperature if it has one, else 0; - V(s) ensemble:
next = min_i V_i(s'); - Q-network:
next = max_a' Q(s', a').
The bootstrap uses the critic's current parameters: whether and how an
agent keeps a target network (SAC does, AVG and PQN do not) is its own
business, as are GAE, Q(lambda) and average-reward targets. Pass
value_target_fn to measure against an agent's exact target.
BiasVorePenalty
dataclass
¶
Bases: Extension
EVarEst critic objective as a critic_loss extension.
Ports the additive form of EVAREST's
:func:make_evarest_loss / :func:make_evarest_ppo_loss /
:func:make_evarest_td_loss. The agent's vanilla critic loss
already contains MSE = bias^2 + Var(residual); this extension
adds the eq. (6) extra penalty term:
penalty(alpha) = ((1 - 2 alpha) / alpha) * Var(q_preds - target_q)
so the total objective becomes
MSE + ((1-2 alpha)/alpha) * Var(residual). alpha=0.5 ⇒ coeff
0 ⇒ identical to vanilla; alpha->0 is the AVEC limit (use a small
alpha rather than zero). coeff_override lets callers fix the
coefficient directly (e.g. to disable the extra penalty entirely
without recomputing alpha). It rebuilds the residual the agent's MSE
term sees from the critic_loss batch: observations,
targets and the critic (critic_params and critic_state,
with actions for a Q-critic; DQN's and PQN's q_state and
actions), and raises on a batch without them rather than add 0.
srank ¶
Kumar et al. (2021) effective rank.
Smallest k whose top-k singular values capture >= 1-delta of the
total spectral mass. feats is the (batch, d) feature matrix.
dormant_fraction ¶
Sokar et al. (2023) dormant-unit fraction.
A unit is dormant when its mean-|activation| score, normalised by
the layer mean, is <= tau. feats is the (batch, d) post-
activation feature matrix.