Training — API reference¶
Every public symbol in skyfall_crl.train, generated from the source.
Generated page — built from the source docstrings, so this file is blank when read on GitHub. Run
mkdocs serveto read it locally, or read the docstrings in the module itself.
What a window and an algorithm are¶
base ¶
What a training algorithm is, and what it is fed.
Three contracts, kept apart because the research code they generalise keeps them together:
Transition/RolloutWindow-- what the collector produces.Policy-- something that chooses an action.Algorithm-- something that consumes windows and improves a policy.
Nothing here imports torch, and nothing here names an environment. A window holds plain Python and NumPy, so the same window describes a CartPole transition and a remediation tool call, and a backend converts to tensors in its own terms. The reference implementation stores token ids directly on the batch, which is why its batch cannot describe an environment whose actions are not text.
The separation of Policy from Algorithm is the one piece with no reference
counterpart -- there, the trainer is the policy, owning the model, tokenizer and optimizer.
That works for exactly one policy family. Keeping them apart costs nothing and means a backend
that owns its model is a case rather than the case.
Transition
dataclass
¶
One environment step, as the harness records it.
Attributes¶
step:
Index within the rollout, counted from the start of the run rather than the
window -- windows are administrative slices of one continuing run.
observation, action, reward:
What the policy saw, what it did, and what it earned. observation and
action are whatever the environment's spaces contain; nothing here inspects
them.
terminated, truncated:
Gymnasium's two ways for a step to be the last one. A persistent environment
never terminates and truncates at an administrative boundary, but that is the
environment's business, not this record's.
regime_id:
The active configuration z_t, for segmenting metrics. Recorded because the
policy is not told it -- inferring the shift is the problem being posed.
boundary:
Whether this step is the first of a new configuration.
reward_components:
Each reward component's weighted contribution, when the environment reports
them. Empty for an environment that supplies only a scalar.
logprob, value:
What the policy thought at the time it acted -- the log-probability it assigned
the action, and its estimate of the state's value. Optional, and absent by
default.
A policy-gradient method needs the log-probability *under the policy that acted*,
and once weights have moved that number cannot be recovered. Recording it here is
the cheap way to keep it. Not every backend wants it: the ported PPO deliberately
re-runs the model over the stored tokens instead, so for that one these stay
unset -- they are here for the backends that do not.
info:
The environment's own info for the step, kept whole so a backend can read
something the harness does not model.
reward_components
class-attribute
instance-attribute
¶
RolloutWindow
dataclass
¶
A fixed-length slice of one continuing run.
A window is an administrative unit -- the span of experience an update is computed from -- and deliberately not an episode. Ending a window does not end an episode, does not reset the environment, and does not make the last step of the window the last step of the episode. That distinction is load-bearing: episode-scoped reward terms fire on the episode's final step, so a window shorter than the episode legitimately sees only the per-step terms, and treating a window edge as an episode edge would fire the episode terms once per window instead of once per episode.
Attributes¶
transitions: The steps, in order. run: Provenance for the whole run -- whatever the environment reported about itself, recorded verbatim and interpreted by nothing here.
Policy ¶
Algorithm ¶
Bases: Protocol
Consumes rollout windows and improves a policy.
Mirrors the four-method seam the reference converged on, with two deliberate changes.
on_regime_shift takes a plain mapping rather than an omegaconf.DictConfig, so
configuration format is not part of the algorithm contract. And policy is exposed,
because a collector needs something to act with and the reference had no way to ask.
on_regime_shift and save_checkpoint are optional in practice: a backend that
neither consolidates on a shift nor checkpoints simply does the nothing they describe.
Optional environment capabilities¶
protocols ¶
What the harness may ask of an environment beyond gymnasium.Env -- all of it optional.
gymnasium.Env is a narrow surface: reset, step, close, the two spaces, and some
metadata. Nothing in it carries run provenance, a configuration label, or a per-step record --
and a trainer wants all three. Typing the harness against the environment that happens to
provide them would hand the coupling the two-layer split exists to prevent to every layer built
afterwards, so instead they live here, are probed rather than required, and are absent without
error. An environment that offers none of it still trains.
Probing goes through has_wrapper_attr / get_wrapper_attr, never .unwrapped.
Gymnasium 1.3 removed Wrapper.__getattr__, so a wrapped environment forwards nothing -- and
.unwrapped is not the fix, because it returns the base environment and therefore skips
any capability a middle wrapper added:
TimeLimit(Described(CartPole)) hasattr(env, "describe") -> False
env.unwrapped has describe -> False
env.has_wrapper_attr("describe") -> True
Reaching for .unwrapped would make the harness quietly report a capability as missing when
it is present, which is worse than failing.
DescribesRuns ¶
Bases: Protocol
Reports what a run was configured with.
The mapping's keys are not specified, deliberately. The one implementation today returns MORPHEUS-shaped keys (a world id, a layout, a reward spec), and pinning those into the contract would make every other environment wrong by construction. The harness records the mapping verbatim as provenance and interprets none of it.
capability ¶
The named attribute of env, looking through any wrappers, or None.
Tries gymnasium's wrapper-aware lookup first, since that is the only one that sees a
capability added partway up a wrapper stack, and falls back to a plain attribute read
for an object that is not a gymnasium.Env at all -- a bare double in a test, say.
describe_run ¶
Run provenance from an environment that offers it, or {}.
{} is a supported answer, not a failure: a Tier-1 environment simply has nothing to
say about itself, and a run that records "nothing was reported" is honest where one that
invents a description is not.
regime_from_info ¶
The configuration label and whether this step begins a new one.
Reads the info["regime"] convention -- {"id": ..., "boundary": ...} -- and
tolerates its absence, a bare string in its place, or a mapping missing either key.
An environment that reports no configuration is stationary as far as the harness is
concerned, which is a legitimate control setting rather than an error.
regime_origin_from_info ¶
What the current configuration repeats, when it repeats one.
A schedule may revisit a configuration under a new label, and the label alone cannot say
that two visits are the same conditions -- which is precisely what forgetting compares.
Read separately from regime_from_info so that function's shape stays as it was for
every environment already reporting the convention without this key.
record_from_info ¶
The environment's own per-step record, if it emitted one.
Only an environment that can decompose its own reward has one to give. When this
returns None the collector builds a minimal record itself, which is what makes the
Tier-1 tier real rather than nominal.
The backend registry¶
registry ¶
Algorithm backends by name, so selecting one is configuration rather than an import.
Selecting an algorithm is a name in a configuration file rather than a different entrypoint, so adding one does not mean forking the harness: register a factory under a name and any config can select it.
Backends resolve on first use, not on import, for a reason particular to this registry:
every real backend imports torch, transformers and friends, and those are an optional extra.
Listing what is available, checking whether a name is taken, and validating a config that names
one must all work in an environment that has none of them installed -- so the table holds
(module, class) strings and nothing is imported until someone asks for the class itself.
The shape is deliberately the same as the reward-component registry, which solved the same problem for the same reason.
The experiment config¶
config ¶
One declarative config that fully specifies a run.
The reference spreads a run across a Hydra tree whose groups are composed at launch, and picks the algorithm by which entrypoint script you invoke -- so "swap the algorithm" means editing Python. Here a run is one document, and the algorithm is a name in it.
Two properties are deliberate and tested:
Reading a config imports nothing heavy. Every backend needs torch; the registry resolves backends lazily and this module only ever holds their names, so a config naming a GPU trainer still loads, validates and prints on a laptop with no ML stack installed.
A backend named here can actually be built. Every shipped backend trains a policy it is handed, so the algorithm section names that policy too -- otherwise "swap the algorithm in the config" would be true only for a backend that owns no model.
A reward is never defaulted. The reward section is optional -- omit it and the
environment's own reward is used, which is the Tier-1 contract -- but naming the section and
leaving the spec out is an error. Every shipped spec describes MORPHEUS's operational model, so
quietly supplying one to an environment that reports something else scores a structural zero on
every step, forever, with nothing in the logs. See RewardConfig.
An unknown key is an error, not a no-op. Every model here rejects keys it does not declare,
so a misspelled or misplaced key -- policy: at the top level instead of under algorithm:,
say -- fails at load naming the key, rather than being silently discarded and leaving the run
subtly different from the document that described it. The free-form surfaces stay free-form by
design: kwargs and params dicts pass through untouched.
EnvironmentConfig ¶
Bases: BaseModel
Which environment to train against.
id is either a registered Gymnasium id (CartPole-v1) or a dotted
package.module:factory for anything not in the registry -- which is how a MORPHEUS
world is selected without this module importing the adapter.
build ¶
Construct the environment. The only place this module imports anything optional.
PolicyConfig ¶
Bases: BaseModel
Which policy the backend trains, as a dotted package.module:factory.
Separate from the algorithm because the two vary independently: the same PPO trains any policy, and the same policy is trained by any of the backends. Resolved at build time, so naming a policy that needs torch costs nothing until the run actually starts.
AlgorithmConfig ¶
Bases: BaseModel
Which backend to train with, by registered name.
The name is checked against the registry on load -- a typo is a load-time error rather than a failure an hour into a run -- and checking it imports nothing, so a config selecting a GPU backend still validates on a machine that could never run it.
The ordering that implies is worth stating: a config naming a backend registered by a plugin must be read after that plugin registers. Built-ins are always available, so this only ever applies to your own.
policy is what makes "select an algorithm in a line of YAML" true for the shipped
backends. Every one of them trains a language policy it is given; without a way to name
that policy, a config could select a backend it could never construct.
RolloutConfig ¶
Bases: BaseModel
How much experience an update sees, and how much the run collects in total.
window_steps is an administrative slice, not an episode: reaching the end of a
window does not end an episode, does not reset the environment, and does not make the last
step of the window the episode's last step. Episode-scoped reward terms therefore fire once
per episode rather than once per window -- so a window shorter than an episode legitimately
sees only the per-step terms.
RegimeConfig ¶
Bases: BaseModel
Which configuration-shift provider drives the run.
Defaults to constant -- a stationary world, which is a legitimate control and not a
continual-learning setting. Say so in the config when that is what you want.
RunConfig ¶
Bases: BaseModel
Identity and outputs. Everything here ends up in the trace's run metadata.
algorithm_id names the family a run belongs to, and it is what a cross-family comparison
groups on. Left unset it is filled in from the algorithm the run selected, so an ordinary
single-algorithm run needs no extra ceremony; state it explicitly when several runs share a
backend but are being compared as different arms -- two configurations of PPO, say.
meta ¶
The run metadata merged into every trace row, omitting what was not set.
ExperimentConfig ¶
Bases: BaseModel
A complete run: an environment, an algorithm, and how much to collect.
reward is optional and its absence is meaningful, not a gap: the environment scores
itself. Supplying it hands scoring to the composable reward system instead. There is no
third state where a reward is chosen for you.
Collecting rollouts¶
rollout ¶
Collecting experience: drive an environment for a fixed number of steps, and record it.
The collector is where the harness meets the environment, and it makes three commitments that matter more than the code that implements them.
A window is an administrative slice of one continuing run, not an episode. It ends when the step budget is spent, and that is all it means. It does not reset the environment, and it does not tell the environment that its episode has ended. Two consequences follow, and both are deliberate:
- Episode-scoped reward terms fire once per episode, so a window shorter than an episode sees only the per-step terms. That is faithful to the reference, whose episode bonus lands on the episode's final step -- and it keeps the composed reward comparable to what the reward system was verified against.
- A configuration shift does not end a window either. If it did, the agent would be handed the boundary for free, and adapting to an unannounced shift is the problem being posed.
Truncation never triggers a reset here. In a persistent environment a reset destroys the world and builds a new one, so auto-resetting at every truncation would throw away the agent's accumulated work and the financial position its reward is computed from. The reference solves this the same way -- its rollout loop contains no reset at all. Resetting is the caller's decision, and for a persistent world the answer is usually "never, after the first".
Every step is recorded, whether or not the environment can describe itself. An environment
that decomposes its own reward puts a record in info["trace"] and we keep it; one that
reports only a scalar gets a record built for it here. That is what makes the Tier-1 contract
real rather than nominal: the metrics library reads traces, so an environment that cannot
produce one would be excluded from measurement no matter what the contract said.
The same rule extends to the configuration label. A world that shifts its own conditions
reports which configuration a step happened under, and that is kept untouched. An environment
that does not -- every Tier-1 environment, since a plain gymnasium.Env has no such concept --
can be handed a configuration provider, and the label is stamped here instead. Without that, such
a run records one undifferentiated segment and adaptation speed, recovery time and forgetting have
nothing to measure: all three are defined relative to a configuration boundary.
RolloutCollector ¶
Drives one environment and yields fixed-length windows of experience.
Parameters¶
env:
Any gymnasium.Env. Nothing beyond that interface is required; anything more it
offers is probed and used if present.
policy:
Anything with act(observation).
window_steps:
How many environment steps make up one window.
writer:
Optional trace sink. Supplied, every step is written as it happens, so a run that dies
halfway still leaves the steps it did take.
auto_reset:
Whether to reset when an episode ends. Off by default, because for a persistent
world a reset is destructive. A bounded environment such as CartPole needs it on to
produce more than one episode.
seed:
Seeds the environment on the first reset only. Deliberately not on subsequent ones:
re-seeding at every auto-reset would start every episode from an identical state, which
looks like reproducibility and is actually a stuck environment.
regime:
Optional configuration provider. Consulted only when the environment reports no
configuration of its own -- a world that shifts its own conditions knows more about them
than a schedule held out here does, so its label always wins. Supplying one is what makes
an environment with no notion of configuration measurable by the six-metric protocol.
step_index
property
¶
Steps taken since the run began -- across windows, not within one.
run
property
¶
Run provenance, as reported by the environment. {} if it reports none.
reset ¶
Begin (or restart) the run. For a persistent world, call this exactly once.
windows ¶
Successive windows until total_steps have been taken.
A final partial window is yielded rather than padded or dropped: the steps happened, and discarding them would misreport the run.
Running an experiment¶
experiment ¶
Running one experiment, from a configuration document to a recorded trace.
A configuration says what to run; this says how a run happens. Keeping the two apart is what lets an experiment be selected rather than written -- but the how still has to exist somewhere, and until this module it existed only in the tests that needed it, copied.
Three things happen here that happen nowhere else, and each is the reason the module exists rather than a convenience it offers:
The algorithm is told when a configuration ends. on_regime_shift is the only signal a
continual-learning backend gets that one set of conditions has given way to another; EWC
consolidates on it and LCM resets context on it. Nothing in the shipped path called it, so a
config-driven EWC run behaved as plain PPO while reporting EWC metrics.
The configuration provider named in the config is actually built and used. It reaches the collector, which stamps the label when the environment reports none -- without which a run on any environment that has no notion of configurations records one undifferentiated segment, and adaptation speed, recovery time and forgetting have nothing to measure.
The environment is closed even when the run fails. For a persistent world that is not housekeeping: an environment left open holds a world, and a world left running costs money.
RunResult
dataclass
¶
What one run produced.
Deliberately not the metrics themselves. Those are computed from the trace, by a library that never sees the environment or the model -- which is what lets a run be scored on another machine, months later, by someone who has neither.
window_metrics
class-attribute
instance-attribute
¶
provenance
class-attribute
instance-attribute
¶
stationary
property
¶
Whether the run saw only one set of conditions.
A legitimate control, and a useless continual-learning benchmark: three of the six metrics are defined relative to a configuration boundary, so a stationary run reports nothing for any of them.
trace_path_for ¶
Where this run's trace belongs, or None if it is not being recorded.
output_dir gives each run its own directory holding traces.jsonl, which is the
layout the metrics library reads directly: point it at the directory a sweep wrote and every
run in it is found, with no manifest saying what is there.
run_experiment ¶
run_experiment(
config: ExperimentConfig,
*,
trace_path: str | Path | None = None,
env: Any | None = None,
algorithm: Any | None = None,
export_dir: str | Path | None = None,
export_external: str | Path | None = None,
) -> RunResult
Run one experiment to its step budget and return what it produced.
export_dir writes an export bundle when the run finishes, before the environment is
closed -- which is the point of exporting from here at all. The observation and action
schemas can only be read while the environment is alive, and this is the one place it and the
trained algorithm are both in hand. export_external keeps the weights outside the bundle;
see export_bundle.
env and algorithm override what the configuration would build. They exist for the
caller that already holds one -- a sweep reusing an expensive policy across seeds, a test
driving a double -- and neither is built here when supplied. An environment passed in is
still closed on the way out, because the alternative is a rule about ownership that is
easy to state and easy to forget.
A matrix of experiments¶
sweep ¶
A matrix of experiments, written once.
A continual-learning result is a comparison: several algorithms, several seeds, the same conditions. Writing that as one configuration document per cell means twenty near-identical files that drift apart, and the reference wrote it as a shell loop over hardcoded checkpoint paths -- which is why its runners are not reusable by anyone else.
Here a sweep is a document like any other. It names what every cell shares, what varies, and the
directory the cells write into; expanding it produces ordinary
ExperimentConfig objects, so a cell is nothing special and can
be run, inspected or printed exactly like a single experiment.
Three properties are deliberate:
Every cell is identified in its own trace. The family, the seed and the task go into the run metadata, not just into a filename, because a filename is not carried when the trace is read back and a cross-family table has to group on something.
Expansion is separate from execution. A matrix can be worth hundreds of GPU-hours, so being able to see exactly what it would run -- before it runs -- is not a convenience.
A cell's seed can reach anything, and the author says what. run.seed seeds the environment
and nothing else -- not the policy's initialisation, not a replay buffer's draw, not a Fisher
estimate's. A sweep whose seeds vary only the world reports a spread across seeds that is zero by
construction, which reads as a strong repeatable result and means the seed reached nothing. Writing
SEED where a seed belongs fills it in per cell. It is a substitution rather than something
inferred from a factory's signature, because the knob is spelled differently by every backend
(seed, replay_seed, fisher_seed) and guessing which one was meant is how a sweep comes
to look seeded without being.
SweepConfig ¶
Bases: BaseModel
What every cell shares, and what varies across them.
families maps a name to the overrides that make that arm what it is -- usually just the
algorithm, but anything a configuration holds. The name becomes the run's algorithm_id,
so two arms may share a backend and still be compared as different families: two
configurations of PPO, say.
tasks may be a list of names or a mapping of name to overrides, because a task is
sometimes only a label and sometimes a different environment.
merge ¶
overlay over base, one level into each section.
Section-wise rather than wholesale: a family that changes the algorithm should not have to restate the environment, and a task that changes an environment keyword should not have to restate the algorithm. Deeper than that is deliberately not merged -- a nested merge makes it impossible to remove a key, and a sweep that cannot express "this arm has no reward" is worse than one that asks for a section to be written out.
expand_sweep ¶
Turn a sweep into the experiments it describes.
Order is family, then task, then seed -- families outermost because that is the axis a partial sweep is most often cut along, so an interrupted run leaves whole arms finished rather than every arm half-done.
Effective rank¶
plasticity ¶
Effective rank: how many independent directions a policy's representation still uses.
Plasticity loss is the failure mode continual learning is for. A policy that has been trained across several configurations tends to collapse its representation into a lower-dimensional subspace, and once collapsed it cannot fit the next configuration however much reward it is shown. The symptom is a monotonic decline in the number of directions the representation actually spans, which is what this measures.
ER(F) = exp(H(p)) where p_i = sigma_i / sum(sigma)
sigma are the singular values of an activation matrix F -- the policy's penultimate-layer
activations across a batch of tokens -- and H is Shannon entropy. The exponential of the
entropy is a continuous rank: 1 when all the spectral mass sits on one component, min(F.shape)
when the spectrum is flat. stable_rank is the integer-valued companion.
Ported from the reference implementation, arithmetic unchanged, including the zero-probability
filter that stops 0 * log 0 producing NaN.
Deliberately torch-free. The maths is a singular-value decomposition, and the six-metric library that consumes it installs without the ML stack. So the computation lives here on NumPy and the capture -- which does need a model -- lives with the policy, which already has one.
effective_rank_entropy ¶
The entropy-based effective rank of a (tokens, hidden) activation matrix.
Returns None rather than raising for anything that has no rank to report -- an empty
matrix, a one-dimensional one, an all-zero one, or a decomposition that does not converge.
A missing measurement is a missing measurement; it is not a reason to end a run, and the
trace records it as None.
stable_rank ¶
The smallest k whose top-k singular values hold threshold of the spectrum.
Integer where effective_rank_entropy is continuous, and coarser for it: it answers
"how many directions carry the mass" rather than "how evenly is the mass spread".
Algorithm backends — advantage¶
advantage ¶
Generalised Advantage Estimation, and how a rollout window becomes episode segments.
Deliberately torch-free. GAE is the subtlest arithmetic in PPO and the easiest place to introduce a silent difference, so it lives where every interpreter can import it and every CI run exercises it -- rather than behind an optional ML extra where it would be tested only on a machine that happens to have one.
In the default configuration, advantage estimation is degenerate, and this is worth knowing
before tuning anything. Episode boundaries reach it only through episode_lengths. Without
them every sample becomes its own one-step episode, the estimate collapses to the one-step
residual r - V, and gamma and lam have no effect whatever their values.
A rollout window carries real terminated / truncated flags, so the episodes can be
segmented properly. Whether they are is a choice: GAEMode selects it. degenerate is
the default because it reproduces the published baselines; episodic uses the window's own
terminal flags and makes gamma and lam live.
GAEMode ¶
Bases: str, Enum
Whether advantage estimation sees episode structure.
DEGENERATE reproduces the reference's vanilla PPO: every sample is its own one-step
episode, so advantage = reward - value and the discount factors are inert.
EPISODIC uses the window's own terminal flags, which is what GAE is for.
episode_lengths_from ¶
Run-length-encode terminal flags into episode lengths.
A window is a slice of a continuing run, so its last episode is usually unfinished. That
trailing run counts as an episode here: its final step bootstraps from zero, which is
exactly what the reference does at every boundary. Treating it as terminal slightly
understates the value of a truncated episode -- the honest alternative would bootstrap from
V(s_T), which the reference never does, so neither do we.
compute_gae ¶
compute_gae(
rewards: Sequence[float],
values: Sequence[float],
episode_lengths: Sequence[int] | None,
gamma: float,
lam: float,
) -> tuple[list[float], list[float]]
Advantages and returns, per episode.
A port of the reference implementation, arithmetic unchanged. Each episode is treated as
terminating, so the value bootstraps from zero at the boundary. When episode_lengths is
missing -- or does not add up to the number of samples -- every sample becomes a one-step
episode and the result reduces to r - V.
Lengths that do not add up to the number of samples are ignored rather than rejected: the
result degenerates to r - V instead of raising. That is forgiving inside a training loop,
and it is the first thing to check when a run's advantages look flat.
Algorithm backends — ewc¶
ewc ¶
Elastic Weight Consolidation over PPO.
The idea is one formula: hold parameters near where they were at the end of the previous configuration, weighted by how much each one mattered there.
L_EWC = ½ · λ · Σ F ⊙ (θ − θ*)²
F is the empirical Fisher -- the squared gradient of the log-likelihood of the tokens the
policy actually generated, averaged over samples, then clipped and normalised per tensor.
Which samples, though, is not "all of them". The draw is stratified by reward band --
solved (r ≥ 0.9) 40%, positive progress 30%, neutral 20%, negative 10% -- and any remaining
budget is filled from the whole pool at random. With fewer samples than budget, which is the
normal case, that fill duplicates: a sample selected twice contributes twice to the average.
The consequence is the point: the Fisher is reward-weighted. EWC protects the weights that
mattered for the episodes that went well, not for experience at large. Set fisher_sample_mix
to change the bands, and fisher_seed to make the draw reproducible.
Two accumulation modes, and they are not equivalent. online keeps one running estimate,
F ← γ·F + F_new with no renormalisation, and replaces the anchor wholly with the newest
snapshot. That γ is fisher_gamma, and the rename is not cosmetic: the reference reads it from
a nested ewc config section, so it never meets PPO's discount factor. Flattening both into
one keyword argument made gamma mean whichever the caller happened to think it meant -- and
set PPO's discount to its default no matter what was passed.
multi_snapshot keeps a FIFO of pairs and sums their fully scaled penalties -- so three
snapshots means an effective 3λ, which is worth knowing before tuning λ.
A second is not reproduced because it cannot be: the noise. The reference puts the model in
training mode for the Fisher pass and then walks its modules switching every nn.Dropout to
eval, commented "for deterministic Fisher estimation". It does not achieve that. Attention
dropout under SDPA is not an nn.Dropout module -- it is a dropout_p argument gated on the
attention module's own training flag -- so it stays active, and two Fisher estimates over
identical inputs differ by more than the values themselves. Measured on a two-layer model: a
maximum absolute difference of 0.75 between consecutive runs of the same call.
That makes the reference's Fisher an irreproducible draw, and there is nothing to be faithful
to in a number that cannot be produced twice. So EWCBackend runs the pass in eval
mode by default, which reaches the reference's stated intent. fisher_dropout="reference"
restores the mode handling as written, noise included.
One reference behaviour is deliberately not reproduced: the silence. Fisher estimation returns
an empty result when the parameter filter matches nothing -- which happens whenever
parameter_scope="lora_only" meets a model without LoRA adapters, i.e. any full fine-tune. In
the reference nothing logs it. The empty result empties the regularizer, has_snapshots stays
false forever, consolidation is never retried, and the run completes as plain PPO while its
log line reports a healthy-looking sample count. A second path does the same thing more quietly
still: an all-zeros Fisher registers as present and contributes a permanently zero penalty.
Both now raise. This is the same failure shape as a reward spec that could read nothing from its environment, and it gets the same treatment: say so rather than return zeros.
DEFAULT_SAMPLE_MIX
module-attribute
¶
DEFAULT_SAMPLE_MIX: Mapping[str, float] = {
"solved": 0.4,
"positive_progress": 0.3,
"valid_neutral": 0.2,
"random": 0.1,
}
FisherDropout ¶
Bases: str, Enum
Whether the Fisher pass is actually free of dropout.
DISABLED puts the model in eval mode, which is what the reference's comment says it
wants. REFERENCE reproduces what it does instead: training mode with nn.Dropout
modules switched off, leaving SDPA's attention dropout running.
EWCMode ¶
ParameterScope ¶
Bases: str, Enum
Which parameters the penalty covers.
All three match on a naming convention -- PEFT's lora_ prefix and TRL's v_head.
A model using neither matches nothing, and EWC would then train as plain PPO while still
reporting a healthy sample count. That case raises instead.
EWCBackend ¶
Bases: PPOBackend
PPO with a Fisher-weighted anchor to the previous configuration's weights.
Everything PPO does is inherited unchanged; the only addition is a penalty applied through
_apply_auxiliary_regularization, which the reference also uses as its extension point
so that the penalty lands inside the clipped gradient norm.
on_regime_shift ¶
Consolidate what the ending configuration taught, before the next one starts.
consolidate ¶
Estimate Fisher on recent experience and anchor the current weights.
estimate_fisher ¶
The empirical Fisher over recent experience.
Raises rather than returning an empty result. In the reference this returns {} with
nothing logged, and the run silently becomes plain PPO -- see the module docstring.
select_fisher_samples ¶
select_fisher_samples(
pairs: Sequence[tuple[Any, Any]],
rewards: Sequence[float],
*,
max_samples: int,
mix: Mapping[str, float] = DEFAULT_SAMPLE_MIX,
rng: Random | None = None,
) -> list[tuple[Any, Any]]
The reference's stratified draw, reproduced including its duplication.
Each band gets round(max_samples * ratio) draws (at least one), and whatever budget is
left is filled from the whole pool. When the pool is smaller than the budget -- the usual
case -- the fill re-draws samples already selected, so they count more than once in the
average. That is what makes the estimate reward-weighted rather than uniform.
build_parameter_filter ¶
A predicate selecting the parameters EWC protects.
Algorithm backends — her¶
her ¶
Hindsight replay, as the reference implements it.
This is not hindsight experience replay in the usual sense, and the docstring says so because the name does not. There is no goal representation, no goal-conditioned reward function, and no achieved-goal extraction from state. The entire mechanic is one line -- the terminal step's reward is clamped up to a constant:
out[-1] = max(out[-1], success_reward)
The future strategy adds a stochastic version of the same clamp to earlier steps, with
probability 1/(m-t) for step t of an episode of length m. The "future goal" it samples
is never read; it only decides whether to apply the same constant. So future is final
plus noise, not a different mechanism.
Ported faithfully anyway, because the runs this package is measured against were produced by it.
What is not inherited silently is the reason success_reward defaults to 0.5: that is
MORPHEUS's theoretical reward ceiling, R_ub = w_l + w_p = 0.25 + 0.25. Against a reward on any
other scale it is an arbitrary number, and the admission threshold is tuned to the same scale.
The relabelling is pure arithmetic on floats, so it lives here without torch and is checked against the reference by a differential test.
RelabelStrategy ¶
HERBuffer ¶
Stores relabelled transitions and replays them into later batches.
Parameters¶
capacity:
-1 disables replay entirely. Note the reference's asymmetry, reproduced here:
0 is enabled with room for a single entry, not disabled.
admission_threshold:
A transition enters only when its relabelled reward is strictly greater. The default
0.0 therefore rejects a zero-reward step, and is calibrated against MORPHEUS's
composite reward -- specifically to reject the small exploration bonus a no-op step earns.
success_reward:
The constant the terminal step is clamped up to. 0.5 is MORPHEUS's R_ub; see the
module docstring before reusing it against a different reward.
add ¶
Relabel a batch's episodes and store whatever clears admission.
Episode structure comes from episode_lengths; when it is missing or inconsistent the
whole batch collapses into one episode, so only its last row gets the clamp. That is
the reference's fallback and it fires constantly, because a replay-mixed batch has no
episode lengths to begin with.
sample ¶
Up to count stored transitions, without replacement.
relabel_episode ¶
relabel_episode(
rewards: Sequence[float],
*,
success_reward: float,
strategy: RelabelStrategy | str = RelabelStrategy.FINAL,
rng: Random | None = None,
) -> list[float]
Rewrite one episode's rewards under a substitute goal.
The originals are never mutated -- a copy is returned, which matters because the caller keeps using the real rewards for the live half of the batch.
Algorithm backends — lcm¶
lcm ¶
Latent Context Model over PPO: infer which configuration you are in, and condition on it.
An LSTM reads a sliding window of recent state features and emits a latent vector z plus
regime-class logits. z is trained online by an auxiliary loss that penalises it drifting
within an episode -- the assumption being that the configuration is stable inside an episode and
changes between them.
This is a proxy implementation and the port keeps it that way. Every one of the following is the reference's actual behaviour, reproduced deliberately, because a family that has been changed cannot be compared against the runs it produced:
- The reward channel is dead. The encoder takes
state_dim + 1inputs where the+1is a reward -- and it is hard-zeroed at rollout, in the update, and in the offline warm-up. Every code path in the reference feeds it zero. - The state features are value-blind. They are a bag of key paths: two states with entirely different numbers under the same keys produce an identical vector. Lists are flattened without indices, so one item and a hundred look the same.
- 3,840 of 4,096 dimensions are permanently zero. Only the first 256 are ever written. The
4,096 is a placeholder for an LLM hidden size that was never wired in, and it costs a
Linear(4097 → 256)whereLinear(257 → 256)would do. - The regime head starts untrained, and its softmax feeds a natural-language prefix injected into the prompt. The shipped warm-up checkpoint was trained when the regime list had three entries rather than six, so those weights are shape-dropped on load.
- Consistency loss is not squared, despite its docstring saying so: it is
mean(1 − cos).
The half that makes it LCM rather than PPO-with-an-extra-loss is the prompt prefix, and it
lives outside this class. In the reference the encoder runs at rollout, and
regime_description_from_z turns the regime posterior into a [REGIME: financial_pressure(0.68)
| edi_spike(0.19)] string prepended to the policy's user turn -- that string is the only way the
latent ever reaches the policy. Here the policy is a separate object from the backend, so this
module exposes the prefix (LCMBackend.regime_prefix) instead of injecting it. Nothing
in this package injects it for you: until a harness wires observe and regime_prefix into
whatever builds the prompt, z is computed, trained, and discarded, and this backend is PPO
plus an auxiliary loss. The consistency loss reads only z, so regime_head receives no
gradient online -- that part is the reference, which trains the head in the offline warm-up
alone.
Four things are made explicit rather than inherited, because inheriting them would break:
- The hash is stable by default. The reference uses Python's
hash(), which is salted per process, so its feature basis differs between the warm-up and the rollout and between any two runs. An irreproducible basis cannot be compared to anything.hash="python"restores it. - The warm-up checkpoint is a config path. The reference resolves it from
__file__up three directories, which breaks the moment the module is vendored -- and the asset is not in the repository at all, so that path never resolves anyway. - One
K. The reference declares the window length independently in three places, which must agree or the rollout and the update slide different windows. - The encoder is passed in, not assigned onto a private attribute of another object.
CANONICAL_REGIMES
module-attribute
¶
CANONICAL_REGIMES = (
"baseline",
"edi_spike",
"infra_degradation",
"financial_pressure",
"workforce_constrained",
"baseline_revisit",
)
HashMode ¶
Bases: str, Enum
Which hash buckets the state features.
STABLE is deterministic across processes. PYTHON reproduces the reference, whose
features are salted by PYTHONHASHSEED and therefore differ between runs.
RegimeContextEncoder ¶
LSTM over a window of state features; emits a unit-norm latent and regime logits.
lstm
instance-attribute
¶
LCMBackend ¶
Bases: PPOBackend
PPO plus an online-trained latent context encoder.
Defaults to gae="episodic" because in the reference LCM is the only configuration that
populates episode lengths -- so it is the only one whose GAE was ever multi-step. Selecting it
here reproduces that rather than quietly running the degenerate form.
observe ¶
Record one observation for the rolling window the encoder reads.
Called at acting time, not update time: the reference slides its window across the rollout, so the prefix a step sees reflects the steps before it.
encode_recent ¶
The latent and regime logits for the window observed so far.
Left-padded with zeros to a full window, as the reference pads, so a run's first steps produce a prefix rather than an error.
regime_prefix ¶
The [REGIME: ...] string to prepend to the policy's prompt.
Prepending it means editing a prompt, and the backend does not own the prompt -- so this returns the prefix and the caller decides where it goes.
load_warmup ¶
Load a pre-trained encoder, dropping tensors whose shape no longer matches.
strict=False alone does not silence a size mismatch -- it ignores missing and
unexpected keys only -- so mismatched shapes are filtered explicitly. The regime head is
the usual casualty, and it is the one feeding the prompt prefix.
state_features ¶
state_features(
state: Any,
*,
hash_mode: HashMode | str = HashMode.STABLE,
width: int = STATE_DIM,
) -> list[float]
A bag-of-key-paths feature vector for one observation.
Values are traversed but never read -- only the set of dotted key paths matters. That is the reference's featurisation and it is why two structurally identical states with different numbers are indistinguishable here.
regime_description ¶
The natural-language regime prefix, from the head's logits.
Reproduces the reference's [REGIME: name(p) | name(p)] format exactly, because it is a
prompt fragment: change the wording and a policy trained against the old one is being fed
something it has never seen.
The probabilities are near-uniform unless a warm-up checkpoint was loaded, since nothing trains the head online. That is worth knowing before reading meaning into them.
consistency_loss ¶
Penalise the latent drifting between adjacent steps of one episode.
mean(1 − cos). The reference's docstring says "mean squared cosine distance"; its
arithmetic does not square, and the arithmetic is what was run.
Algorithm backends — policy¶
policy ¶
A language policy: generates a response, and remembers the tokens it generated.
PPO needs the exact tokens it is updating over, and the harness's rollout window does not carry them -- deliberately, since no other kind of environment has any. So the policy records what it generated, and the backend that owns it collects the record. The reference reaches the same place by having the trainer be the policy; keeping them separate costs one small buffer and keeps the harness substrate-neutral.
The value head is TRL's when TRL is installed, and a plain linear head otherwise. They are the
same thing -- a Linear(hidden, 1) over the last hidden state -- so the fallback is not a
degradation, it is the same computation without the dependency.
ValueHead ¶
A scalar value estimate from the last hidden state.
Exists so the backend does not require TRL. TRL's AutoModelForCausalLMWithValueHead
wraps the same Linear(hidden, 1); using ours changes the initialisation, not the model.
LanguagePolicy ¶
Wraps a causal language model so it can act, be scored, and be trained.
Parameters¶
capture_effective_rank:
Measure the representation's effective rank on every action. Off by default -- it is a
second forward pass plus a decomposition per step, and it is a diagnostic rather than
something the update needs.
model, tokenizer:
Injected rather than built from a name, so a policy can be constructed over any model
you already hold -- including a randomly initialised one, which is what makes a run
testable without a download. language_policy builds one from a name or a
checkpoint path when that is what you want.
json_open_ids
instance-attribute
¶
json_open_ids: set[int] = {
token_id
for token, token_id in getattr(
tokenizer, "get_vocab", dict
)().items()
if isinstance(token, str) and token.startswith("{")
}
act ¶
Generate a response and parse it into an action, recording the tokens.
parse
staticmethod
¶
The action inside a response, or the raw text when there is no JSON object in it.
Returning the text rather than raising is deliberate: a malformed generation is an invalid action for the environment to refuse, not a reason to end training.
take_generated ¶
The tokens generated since the last call, and clears the record.
Cleared on read so a window can never be scored against a previous window's tokens -- the failure that would pair rewards with the wrong actions and be nearly invisible.
forward ¶
Logits over the sequence, and the value at its final token.
save ¶
Write the model, the tokenizer, and the value head.
The value head is saved separately because it is not part of the language model, and
save_pretrained therefore knows nothing about it.
load_value_head ¶
Restore a saved value head. Returns whether one was found.
Worth doing explicitly, because the failure is invisible: a resumed run whose critic
was not restored continues the policy from where it left off while starting the
critic from random initialisation, so its advantages are meaningless until the critic
re-converges, and no metric reports it. loaded_value_head records which happened.
Four filenames are accepted, in both PyTorch and safetensors form, because upstream
training stages do not agree on one: a value-warm-start stage writes v_head.pt where
a PPO checkpoint writes v_head.safetensors.
from_checkpoint
classmethod
¶
Rebuild a policy from a saved checkpoint, value head included.
This is what makes SFT-then-PPO a chain rather than two unrelated stages: the
supervised stage's weights and its critic both carry forward. Accepts either a
full-weights checkpoint or a LoRA adapter -- see _load_pretrained.
language_policy ¶
Build a policy from a model name or path -- the config-facing entry point.
A checkpoint directory carrying a saved value head is restored through
LanguagePolicy.from_checkpoint, so an SFT stage's critic carries into PPO rather
than being silently discarded. Anything else is loaded fresh.
Algorithm backends — ppo¶
ppo ¶
PPO for a language policy, ported from the reference implementation.
The algorithm is reproduced as the reference computes it, including the parts that look like defects, because the runs this package is measured against were produced by exactly this arithmetic. Every one of them is named, defaulted to the reference behaviour, and switchable:
============================ ==========================================================
Option Reference behaviour (the default)
============================ ==========================================================
gae degenerate -- the reference's collector supplies episode
lengths only when a latent-context encoder is attached, so
vanilla PPO reduces to advantage = reward - value and
gamma/lam do nothing. episodic uses the window's
own terminal flags, which we have and it did not.
ratio sequence_mean -- the ratio is built from the mean
per-token log-probability, so it is the geometric mean of the
per-token ratios. clip_range therefore constrains average
drift, not each token, and is far looser at sequence level
than textbook PPO. per_token is the textbook form.
action_tokens qwen_think -- gradient flows through the whole response
for a thinking model, and from the last JSON-opening token
otherwise. all reinforces everything.
old_policy_pass train -- the reference leaves the model in training mode
for the no-grad behaviour pass, so dropout is active and the
ratio is not 1.0 even on the first minibatch. This looks
unintended. eval disables dropout for that pass.
============================ ==========================================================
Two further reference behaviours are reproduced without a switch, because they are not
parameters: there is no entropy bonus, and kl_coeff does not weight a KL divergence --
it weights an L2 anchor to the weights the run started from, applied once per optimizer step.
The genuine KL is computed only as a diagnostic, and under LoRA there is no reference model, so
it is always 0.0. The anchor is off by default (kl_coeff=0.0), which is also the point at
which the initial snapshot is skipped entirely rather than held for a run that will never use it.
Old log-probabilities and values are recomputed from the stored tokens on every update
rather than captured at generation. That is the reference's choice and is kept; it is also why
Transition.logprob goes unused here.
Hindsight replay lives here rather than in a backend of its own, because in the reference it
is not an algorithm -- it is a buffer switched on by replay_size != -1, and its shell script
runs the same entrypoint as vanilla PPO. Set replay_size and this is the reference's "HER".
Two arithmetic guards are not switchable, because without them a run fails silently rather than differently:
- The advantage normalisation is guarded (
numel() > 1), as the reference guards it. Our first port dropped the guard, and a one-sample window then wrote NaN into the weights. - The optimizer and the gradient clip cover the policy's parameters, not the language
model's. The reference's value head is a submodule of its wrapped model, so its
model.parameters()already includes the critic; ours is held beside the model, where the same expression silently excludes it.
GAEMode ¶
Bases: str, Enum
Whether advantage estimation sees episode structure.
DEGENERATE reproduces the reference's vanilla PPO: every sample is its own one-step
episode, so advantage = reward - value and the discount factors are inert.
EPISODIC uses the window's own terminal flags, which is what GAE is for.
ActionTokens ¶
Bases: str, Enum
Which part of a response is reinforced.
QWEN_THINK is the reference's rule, and depends on the tokenizer having <think> /
</think> and on the action being JSON. ALL reinforces the whole response, which is
what a model with neither convention needs.
RatioMode ¶
OldPolicyPass ¶
TokenBatch ¶
What PPO actually needs from a window: tokens, rewards, and episode structure.
A RolloutWindow records observations and actions in
whatever form the environment uses. Turning those into tokens is the policy's job -- it
owns the tokenizer and it is the thing that generated them -- so a policy that intends to be
trained by PPO records what it generated, and this collects it.
PPOBackend ¶
Proximal policy optimisation over a language policy.
Parameters¶
policy:
Something that acts and records what it generated -- see TokenBatch. Owning
the policy is how this backend gets tokens; the harness contract deliberately does not
carry them, because no other kind of environment has any.
optimizer:
Optional. Built over trainable_parameters when omitted -- the language model
and the critic.
replay_size:
-1 disables hindsight replay. Anything else turns this backend into the
reference's "HER" configuration; see her for what
that name does and does not mean.
gae, ratio, action_tokens, old_policy_pass:
The four switches described in the module docstring. Defaults reproduce the reference.
replay
property
¶
The hindsight replay buffer. Disabled unless replay_size was set.
trainable_parameters ¶
Every parameter this backend optimises -- the critic included.
The reference reaches the same set through model.parameters() because its value
head is a submodule of a wrapped model. Ours is held beside the model rather than
inside it, so that expression would silently exclude it: the critic would never be
stepped, its gradients would never be zeroed, and every advantage would be computed
against a permanently random baseline. Asking the policy for its parameters is what
makes the two arrangements equivalent.
update ¶
One PPO update over a window of experience.
on_regime_shift ¶
Nothing to consolidate: plain PPO carries no memory across configurations.
That is the point of the baseline -- it is what EWC and LCM are measured against.
batch_from ¶
Collect what the policy generated for this window.
Episode structure comes from the window's own terminal flags, which is more than the
reference ever had -- and is discarded again under the default gae mode, so that a
run reproduces the reference rather than quietly improving on it.
mix_replay ¶
Fold hindsight replay into the batch, exactly as the reference sequences it.
Order matters and is not the obvious one: the fresh batch is added to the buffer
before the sample is drawn, so a replayed row can be a relabelled copy of a row
already in this same batch. The first cycle only fills the buffer -- mixing begins on
the second. And the mixed batch carries no episode lengths, so any run with replay
enabled falls back to degenerate advantages from its second update onwards, whatever
gae is set to. All three are the reference's behaviour, reproduced.
Algorithm backends — tokens¶
tokens ¶
Which response tokens carry the policy gradient.
A generated response is not all action. A model that thinks before it acts emits a reasoning block and then the action; a model that explains itself emits prose and then the action. Which of those tokens the gradient flows through is a real choice, and the reference makes a specific one that is easy to miss and hard to reverse-engineer from behaviour.
It is also load-bearing beyond which tokens get gradient. The log-probability that enters the ratio is the mean over the selected slice, so changing the slice changes the divisor and rescales the whole surrogate. A thinking response of 500 tokens and a bare-JSON response of 20 therefore get systematically different per-token gradient magnitudes.
The one property that must hold whatever the mode: the old-policy pass and the new-policy pass have to slice at the same index, or the ratio compares different token sets and means nothing. That is why this is a pure function of the ids -- it cannot drift between the two passes.
Torch-free apart from an annotation: the selection is an index computation over token ids, so it is tested on every interpreter and the caller applies the resulting slice.
ActionTokens ¶
Bases: str, Enum
Which part of a response is reinforced.
QWEN_THINK is the reference's rule, and depends on the tokenizer having <think> /
</think> and on the action being JSON. ALL reinforces the whole response, which is
what a model with neither convention needs.
action_token_start ¶
action_token_start(
ids: Sequence[int],
*,
mode: ActionTokens | str = ActionTokens.QWEN_THINK,
json_open_ids: Collection[int] = (),
think_open_id: int | None = None,
think_close_id: int | None = None,
) -> tuple[int, bool]
Where the reinforced slice begins, and whether this response was a thinking one.
Three cases in priority order, reproducing the reference:
- A thinking response --
</think>appears anywhere, or<think>is the first token -- reinforces the whole response. Excluding the reasoning tokens gives them zero gradient, and the model stops producing think blocks within a few updates. - Prose before the action reinforces from the last JSON-opening token, so context-specific preamble is not amplified.
- Neither reinforces the whole response.
</think> is searched for anywhere rather than only at position 0 because a model
commonly emits a newline first.