APG¶

Analytic Policy Gradient trains a deterministic controller by differentiating the closed-loop return through the simulator itself. Each update runs the controller for horizon steps on a batch of systems and follows the exact gradient of the summed reward with respect to the controller's parameters, back through every step of the physics. There is no critic, no replay buffer and no exploration noise, so it needs an environment whose transitions carry gradients.

Papers: APG is the general method of training a policy by backpropagation through a differentiable simulator. APG.contextual_controller is the configuration of Busetto, Breschi, Forgione, Piga and Formentin, One controller to rule them all, 2024 (arXiv:2411.06482): a transformer over the tracking history closed by a PID layer, trained across a family of systems.

Can I use?¶

Feature	APG
Discrete actions	no
Continuous actions	yes, tanh-bounded by default (`squash=True`), mapped onto the environment's bounds
Memory (`memory=`)	yes: GRU, LSTM, transformer, Mamba
Environments	gymnax environments with transition gradients: Pendulum, MountainCarContinuous, PointRobot, Reacher, Swimmer
Families of systems	yes, `system_class=`
Several seeds in one call	yes, `train(seed=[...])`
Extension phases	`pretrain`, `post_update`, `eval_metrics`, `actor_loss`
Resume and checkpoints	yes; a split run equals the uninterrupted one (probe, and the resume contract with memory)

Brax and Playground environments are not supported: they take no per-step parameters and have their own gradient path. An environment that cannot expose its gradients raises instead of training on zero gradients.

Example¶

Train one controller for a family of pendulums whose mass and length vary:

import gymnax
from ajax import APG
from ajax.environments.system_class import UniformPerturbation

_, params = gymnax.make("Pendulum-v1")
systems = UniformPerturbation(params, fields=("m", "l"), scale=0.1)  # +-10 %
agent = APG("Pendulum-v1", n_envs=16, horizon=50, system_class=systems)
state, aux = agent.train(seed=[1, 2], n_timesteps=80_000)
print(aux.matching_loss.shape)     # (seeds, updates): 80,000 / (16 * 50) = 100 updates
print(aux.matching_loss[:, ::20])  # minus the closed-loop return

train returns (state, aux) whatever the logging: aux holds every update's loss, matching_loss (minus the mean closed-loop return), m_rmse, grad_norm (the gradient's global norm, before clipping) and timestep. With a logging_config, every log_frequency steps it also evaluates the controller on num_episode_test systems, drawn afresh at each evaluation, and sends Eval/episodic mean reward and Eval/m_rmse to the logging backend. aux.evaluations then holds what each update logged, NaN on the updates that did not evaluate, with or without a backend.

The in-context controller of Busetto et al., trained in stages on a reference-tracking task (train_curriculum, the paper's Algorithm 2):

# doc: skip (compiles a transformer and trains 700,000 steps; tests/agents/APG/test_APG.py covers the curriculum)
import gymnax
import jax.numpy as jnp
from ajax.agents.APG import APG, CurriculumStage, train_curriculum
from ajax.environments.model_reference import LinearReferenceModel, ModelReferenceWrapper, StepReference
from ajax.environments.system_class import FixedSystem, UniformPerturbation
from ajax.wrappers import InitialStateWrapper

plant, params = gymnax.make("Pendulum-v1")
fixed_start = lambda key, s, p: s.replace(theta=jnp.asarray(0.0), theta_dot=jnp.asarray(0.0))
plant = InitialStateWrapper(plant, fixed_start)  # the paper keeps the initial conditions fixed
task = ModelReferenceWrapper(  # track a reference model's step response
    plant,
    StepReference(horizon=100, min_value=-0.5, max_value=0.5, min_duration=20, max_duration=50),
    LinearReferenceModel.first_order(a=0.9, b=0.1, c=1.0, d=0.0),
    output_fn=lambda obs: jnp.arctan2(obs[1], obs[0]),
)
systems = UniformPerturbation(params, fields=("m", "l"), scale=0.05)
stages = [  # the nominal pendulum first, then the whole family
    CurriculumStage(APG.contextual_controller(task, FixedSystem(params), env_params=params), n_timesteps=200_000),
    CurriculumStage(APG.contextual_controller(task, systems, env_params=params), n_timesteps=500_000),
]
*_, (state, aux) = train_curriculum(stages, seed=0)  # one (state, aux) per stage: keep the last

Key hyperparameters¶

Argument	Default	What it does
`n_envs`	8	systems drawn and simulated in parallel per update
`horizon`	100	closed-loop steps differentiated through
`learning_rate`	5e-4	AdamW step size
`weight_decay`	0.0	AdamW weight decay (0 is plain Adam)
`beta_1`, `beta_2`	0.9, 0.999	Adam's moment rates (its epsilon is torch's 1e-8)
`lr_schedule`	None	`"warmup_cosine"`: linear warmup over `warmup_steps` updates, then cosine decay to the call's last update; when `train` keeps the optimiser (`reset_optimizer_on_resume=False`), to the run's, set by the call that builds it (Checkpoints)
`warmup_steps`, `lr_end_fraction`	0, 0.1	warmup length, and the final rate as a fraction of `learning_rate`
`max_grad_norm`	1.0	gradient clipping by global norm (`None` turns it off)
`actor_architecture`	`("128", "relu", "128", "relu")`	MLP in front of the memory or output; `()` for a linear embedding
`memory`	None	memory block (`MemoryConfig` or a dict)
`pid`	None	`PIDHeadConfig`: close the controller with a learnable PID layer
`squash`	True	bound actions to [-1, 1] with tanh; without it they are clipped there
`system_class`	None	distribution over the environment's parameters; `None` trains on `env_params` only
`controller_factory`	None	builds a custom controller module (see `Controller` in `ajax.agents.APG.networks`)
`episode_length`, `action_repeat`	None, 1	accepted as by every agent; on gymnax environments, the only ones APG takes, the episode length is the environment's and the repeat must be 1

APG.contextual_controller(env, system_class, ...) sets the paper's values: 8 layers, 4 heads, width 128, context 100, one system per update (n_envs=1), horizon 100, betas (0.9, 0.95) with weight_decay=0 (plain Adam), 5,000 warmup updates then cosine decay to a tenth of the rate, no gradient clipping (max_grad_norm=None), and a PID output layer. Its final_norm=False rebuilds a controller saved before the transformer ended with its final LayerNorm (see Memory). Such a controller's parameters also map onto the new network, which computes the same outputs: controller_params_with_final_norm(params, memory) moves the controller's LayerNorm, which played that final norm's part, into the transformer. It maps the parameters only, so a resumed training state must restart its optimizer.

Notes¶

Budget. One update uses n_envs * horizon steps of n_timesteps.
Actions. Squashed or clipped, the controller's [-1, 1] action is mapped onto the environment's bounds inside the differentiated rollout, so on gymnax Pendulum (torque in [-2, 2]) the controller reaches the whole range.
Deviations from Busetto et al., each backed by runs on the pendulum tracking task (see the APG docstrings): actions are tanh-bounded (the reference leaves them unbounded); gradients are clipped, since BPTT over 100 steps across a family of systems diverged without it; the PID layer starts as the identity (K_P = 1, K_I = K_D = 0) instead of all three gains at 1, which trained markedly worse under the bounded action (fresh-system M-RMSE 0.33 against 0.25). These runs predate the mapping of actions onto the environment's bounds: the pendulum's torque was then capped at 1 of its
Not backed by runs: the transformer's LayerNorms, its final one included, keep a bias and an epsilon of 1e-6, where the reference's have no bias and 1e-5 (buswayne/in-context-controller at 7eb094ce17: transformer_onestep.py:33 and :156, train_onestep_lin_CL.py:103). The paper trains each curriculum stage to a loss tolerance; Ajax gives each stage a fixed n_timesteps.
Initialisation. A transformer memory starts as the reference's GPT: kernels from N(0, 0.02²), the residual projections at 0.02/√(2 n_layers), zero biases (transformer_onestep.py:162-166, :189-195; MemoryConfig(init="nanogpt")). Other memory kinds keep flax's.
Curriculum. train_curriculum resumes each stage from the previous one's parameters and, by default, restarts the optimizer and its schedule at each stage. All stages must share n_envs, the network and the PID layer.
Tasks. ajax.environments.model_reference turns a plant into a reference-tracking task: ModelReferenceWrapper, a StepReference signal and a LinearReferenceModel whose response the plant must match. On such a task the loss is the paper's closed-loop matching cost and m_rmse its M-RMSE.
Differentiable simulation covers system classes, transition gradients and the reference-tracking wrapper.

Results¶

Unit tests (tests/agents/APG) check its pieces against exact contracts, without learning, including the curriculum. The probes (tests/probing) train it briefly on tiny environments with known answers: it reaches the exact optimum of a small problem whose best action is reachable only through the transition's gradient, and a run split by a checkpoint equals the run left whole. On the CPU speed benchmark (benchmarks/agent_bench.py: Pendulum tracking, 4 environments, 8,000 steps) APG ran at 34,958 steps per second of execution (5,339 with compilation) at main a8a13c2 (PERFORMANCE_REPORT.md).

API¶

ajax.agents.APG.APG.APG ¶

APG(
    env_id,
    n_envs=8,
    horizon=100,
    learning_rate=0.0005,
    weight_decay=0.0,
    beta_1=0.9,
    beta_2=0.999,
    lr_schedule=None,
    warmup_steps=0,
    lr_end_fraction=0.1,
    max_grad_norm=1.0,
    actor_architecture=("128", "relu", "128", "relu"),
    memory=None,
    pid=None,
    squash=True,
    system_class=None,
    env_params=None,
    episode_length=None,
    controller_factory=None,
    extensions=(),
    action_repeat=1,
)

Bases: ActorCritic

Parameters:

Name	Type	Description	Default
`env_id`	`str \| EnvType`	gymnax env id or prebuilt gymnax env (e.g. a `ModelReferenceWrapper`). Must support transition gradients.	required
`n_envs`	`int`	systems sampled (and simulated in parallel) per update.	`8`
`horizon`	`int`	closed-loop rollout length differentiated through.	`100`
`learning_rate`	`float`	AdamW step size.	`0.0005`
`weight_decay`	`float`	AdamW's decoupled weight decay, multiplied by the (scheduled) learning rate (`OptimizerConfig.weight_decay`, every agent's); `0` is plain Adam.	`0.0`
`beta_1`	`float`	AdamW first-moment decay.	`0.9`
`beta_2`	`float`	AdamW second-moment decay.	`0.999`
`lr_schedule`	`Optional[str]`	`None` (constant) or `"warmup_cosine"` (linear warmup over `warmup_steps` updates, cosine decay to `learning_rate * lr_end_fraction` at the call's last update). When `train` keeps the optimizer (`reset_optimizer_on_resume=False`), a resumed call runs the schedule of the call that built it, the run's first or a checkpoint's skeleton: given that flag, that call decays to the run's last update, at its `total_timesteps`.	`None`
`max_grad_norm`	`Optional[float]`	global-norm gradient clipping, on by default: BPTT over 100 steps of a nonlinear plant across a system class diverged without it in our runs (the loss rose 1.2 -> 1.7 M-RMSE within a stage) and was stable with it. `None` disables, as :meth:`contextual_controller` does by default (the reference code does not clip).	`1.0`
`actor_architecture`	`Sequence[str]`	MLP encoder in front of the memory/output; `()` for a linear embedding only (the paper's choice).	`('128', 'relu', '128', 'relu')`
`memory`	`Optional[Union[MemoryConfig, dict]]`	:class:`MemoryConfig` (or dict) for a memory-augmented controller; `None` for a memoryless one.	`None`
`pid`	`Optional[PIDHeadConfig]`	:class:`PIDHeadConfig` to close the controller with a learnable PID layer; `None` for a plain output.	`None`
`squash`	`bool`	tanh-bound actions to `[-1, 1]` (Ajax convention). Squashed or not, `[-1, 1]` is mapped onto the env's bounds (an unsquashed action is clipped there first).	`True`
`system_class`	`Optional[SystemClass]`	distribution over `EnvParams` to train across; `None` trains on the env's own params only.	`None`
`controller_factory`	`Optional[Callable[..., Any]]`	optional builder of a custom policy module, called with `env_args`, `network_args`, `action_dim`, `pid` and `squash` keywords; must honour the :class:`~ajax.agents.APG.networks.Controller` contract (`stateful`, `initialize_carry`, call convention). Lets downstream projects plug their own architectures into the same trainer without touching Ajax.	`None`
`extensions`	`Sequence[Extension]`	composable research features (see :mod:`ajax.extensions`).	`()`

contextual_controller `classmethod` ¶

contextual_controller(
    env_id,
    system_class=None,
    n_envs=1,
    horizon=100,
    n_layers=8,
    n_heads=4,
    d_model=128,
    context=100,
    final_norm=True,
    learning_rate=0.0005,
    weight_decay=0.0,
    warmup_steps=5000,
    pid=_DEFAULT_PID,
    beta_2=0.95,
    lr_schedule="warmup_cosine",
    max_grad_norm=None,
    actor_architecture=(),
    **kwargs,
)

The in-context controller of Busetto et al. 2024.

Defaults follow the paper (Table 2: 8 layers, 4 heads, context 100, width 128, b=1 system per iteration, N=100 horizon) and the reference code (AdamW betas (0.9, 0.95), eps 1e-8, 5k-step warmup then cosine decay to lr/10, no gradient clipping). Deviations, each backed by runs on the pendulum tracking test bed: actions are tanh-bounded (Ajax convention; the reference leaves u unbounded), and the PID head starts at the identity rather than the reference's all-ones coefficients (see _DEFAULT_PID). Not reproduced: the reference (buswayne/in-context-controller at 7eb094ce17) sets its scheduled rate on its first parameter group only (train_onestep_lin_CL.py:241), so its LayerNorm weights keep the peak rate, and decays only tensors of two or more dimensions (transformer_onestep_CL.py:72-77; moot at weight_decay=0). Not backed by runs: the LayerNorms, ln_f included, keep a bias and flax's epsilon 1e-6, where the reference's have no bias and epsilon 1e-5 (buswayne/in-context-controller at 7eb094ce17: transformer_onestep.py:33 calls F.layer_norm with 1e-5, :156 builds ln_f = LayerNorm(n_embd, bias=config.bias), and train_onestep_lin_CL.py:103 defaults --bias to False). The transformer starts as the reference's GPT does (APG_MEMORY_INITS). final_norm=False rebuilds the transformer without its final LayerNorm, only to restore a checkpoint saved before it (see :class:~ajax.networks.memory.MemoryConfig); an expand-then-contract shim (CLAUDE.md, refactoring rule 10), it is removed together with MemoryConfig.final_norm. Extra kwargs go to :class:APG.

ajax.agents.APG.curriculum.CurriculumStage `dataclass` ¶

CurriculumStage(agent, n_timesteps, name='')

ajax.agents.APG.curriculum.train_curriculum ¶

train_curriculum(
    stages,
    seed=42,
    num_episode_test=10,
    logging_config=None,
    reset_optimizer=True,
)

Train stages in order, each resumed from the previous one.

Returns one (agent_state, aux) pair per stage (what train returns). reset_optimizer re-initialises the optimizer and its learning-rate schedule at every stage boundary.

ajax.agents.APG.networks.PIDHeadConfig `dataclass` ¶

PIDHeadConfig(
    use_p=True,
    use_i=True,
    use_d=True,
    kp_init=1.0,
    ki_init=0.0,
    kd_init=0.0,
    anti_windup=None,
)

Hyper-parameters of the controller's PID output layer (the fields of :class:~ajax.modules.pid_head.PIDOutputHead but its width).

ajax.agents.APG.networks.Controller ¶

Bases: Module

ajax.agents.APG.networks.controller_params_with_final_norm ¶

controller_params_with_final_norm(params, memory)

A Controller's parameters saved before its transformer or mamba memory ended with a final LayerNorm, mapped onto the network built now.

The Controller then applied its own LayerNorm (norm) to the memory's output, before the output layer; it now leaves that to the core's final LayerNorm (ln_f). Both are the same LayerNorm at the same place, so moving norm to memory_cell/core/ln_f keeps the outputs identical. params is the variables dict as saved ({"params": ...}), memory the new network's config. Only the parameters map: an optimizer state saved with them keeps the old tree, so a resumed training state must restart its optimizer.

Expand-then-contract (CLAUDE.md, refactoring rule 10): goes with MemoryConfig.final_norm.