Skip to content

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.