Algorithms and Baselines¶
Experiment learners operate independently per agent on PettingZoo parallel
environments. Baselines use discrete actions; parameterized OmSh additionally
uses a categorical/raw-Gaussian focal action. Trainers return a shared
TrainResult containing episode history, metrics, policy adapters, and
agent-specific training state.
| Reported as | Notebook switch | Trainer | Role |
|---|---|---|---|
| IPPO | run_ippo |
IPPO |
Unconstrained baseline |
| IPPO-Lagrangian | run_ippo_lag |
IPPO_Lagrangian |
Constrained-training baseline |
| CPO | run_cpo |
CPO |
Constrained-training baseline |
| OmSh-Monotone | run_ippo_shielded |
IPPO (or ParameterizedIPPO) behind TransitionShield |
Primary method |
| IPPO-TrueShielded | run_ippo_true_shielded |
IPPO behind TruePolicyTransitionShield |
Oracle reference |
IPPO¶
IPPO is the JAX/Flax independent PPO baseline. Each agent owns a separate
actor-critic and optimizer; parameters are not shared. Python collects one live
environment rollout, then JAX batches the PPO updates. The trainer supports
Gymnasium Discrete and arbitrary-rank MultiDiscrete action spaces, including
non-zero starts. It currently supports num_envs == 1 and requires synchronous
episode completion.
IPPO-Lagrangian¶
IPPO_Lagrangian adds one cost critic and dual variable per agent. It reads
infos[agent_id]["cost"], forms separate reward and cost advantages, and uses
the scaled policy advantage (A_reward - lambda * A_cost) / (1 + lambda).
The dual compares cost_limit with completed episodic cost. Partial episodes
persist across rollout boundaries, and no dual update occurs until an episode
completes. This avoids treating an arbitrary rollout fragment as a complete
constraint sample.
CPO and ICPO¶
ICPO is a local PyTorch categorical-action port of constrained policy
optimization. Each agent owns an actor, reward critic, cost critic, and learned
future-failure predictor. Actor updates use conjugate gradient, KL/Fisher trust
regions, constrained case selection, and backtracking line search.
CPO is the plain preset with learned cost shaping disabled, and it is the
constrained baseline the experiment notebooks run (run_cpo, histories in
cpo_histories.pkl). ICPO, which shapes the raw cost with predicted
near-future failure risk, remains available for custom drivers.
Learned Shielded IPPO (OmSh)¶
JAX IPPO acts through TransitionShield. Its focal observation contains the
environment observation plus current budget and opponent-level floor. The
learner samples only currently admissible pure actions: the trainer masks the
logits with the shield's action mask before sampling. External policies can
still be replaced by the highest-reward admissible pure action using
certificates built from the learned environment graph and opponent stack.
ParameterizedIPPO and ParameterizedTransitionShield provide three further
conditions: pairwise primitive-action mixtures with egalitarian budgets, pure
actions with learned successor budgets, and their combination. The trainer
stores the exact categorical and raw-Gaussian augmented action; projection and
primitive sampling belong to the environment transition.
MultiAgentTransitionShield protects several roles in one environment step,
with one controller, opponent model, posterior, floor, and budget per role. With
shield_padding_radius > 0, the same trainers run against the padded
certificate. See Probabilistic Shielding.
True-Policy Shielded IPPO (TrueShielded)¶
TruePolicyTransitionShield uses an exact transition graph. It starts with
uniform opponent policies and can rebuild a pending shield bundle from the live
opponent policies at a configured update interval; pending bundles install at
reset. The interval is true_policy_refresh_interval_updates (every update for
small graphs, every 64 for three-agent Pursuit).
It is an oracle reference: exact dynamics and the opponents' live policies are information no deployed agent has. Comparing it with OmSh separates world- and opponent-model error from the shielding rule itself. It is not evidence for an arbitrary-opponent guarantee.
Boundary Semantics¶
All trainers infer hidden environment episode caps. True terminations stop GAE and do not bootstrap. Truncations stop the recursive trace but retain the actual next-state bootstrap. Shared-environment agents must finish together; asynchronous per-agent endings are rejected.
Checkpoints and Frozen Policies¶
Every trainer saves a schema-3 checkpoint with a restorable inference policy for every agent. OmSh checkpoints also carry the post-training controller state, so a frozen OmSh policy can be evaluated as deployed. See Final-Policy Evaluation.