JAX IPPO Migration¶

The new src/rl/ippo.py and src/rl/ippo_lag.py modules are repo-native JAX trainers that work directly against PettingZoo ParallelEnv environments.

Design¶

• IPPO is fully independent: each agent gets its own Flax module, optimizer state, and update loop. There is no parameter sharing.
• Rollouts are collected in Python from the live PettingZoo env, then batched into JAX arrays for PPO updates.
• Observations are flattened to float32 vectors before entering the policy/value networks.
• Discrete and arbitrary-rank MultiDiscrete action spaces are native. The policy has one categorical head per MultiDiscrete component, sums component log probabilities/entropies for PPO, preserves the space's original action shape, and translates zero-based network indices through Gymnasium start offsets before stepping the environment.
• The current trainer targets num_envs == 1. Multi-env PettingZoo orchestration is intentionally deferred.
• src/rl/parameterized_ippo.py extends the same rollout/update structure to factored categorical and raw-Gaussian augmented actions. It masks unused continuous successor slots out of both rollout and update likelihoods and is the on-policy trainer for parameterized OmSh variants.
• A wrapper may expose action_mask(agent_id) for a Discrete policy. IPPO applies the current mask before sampling, stores it per transition, and uses it again for PPO log-probability and entropy evaluation. After stepping, IPPO compares the proposal with last_executed_actions and rejects any rewrite; this prevents an on-policy update from assigning reward to an action whose logged behavior probability described a different action.

Validation and update sizing¶

• validate_ippo_config(...) rejects unknown keys, non-finite values, invalid ranges, non-positive rollout/update sizes, unsupported activations, malformed hidden-layer widths, and any num_envs value other than 1. train(...) independently requires a positive integer timestep budget.
• Environment validation requires a real PettingZoo ParallelEnv with at least one unique possible-agent identifier. Agent observation and action shapes are inferred before model construction.
• Learning-rate annealing uses the exact number of optimizer applications. This includes a short final rollout and the actual divisor selected when a rollout cannot be divided into the requested number of minibatches.
• tests/test_ippo.py owns direct validation, action-space, probability, GAE, reproducibility, asynchronous-ending, partial-rollout, and learning tests.

Policy Interface¶

• IPPO(cfg, seed=...) and IPPO_Lagrangian(cfg, seed=...) expose .train(env, timesteps) and return a TrainResult.
• TrainResult.policies contains lightweight policy adapters with:
• action_probs(obs) for future shielding / analysis hooks
• act(obs, rng=None) for deterministic or sampled action selection outside training

These adapters are deliberately smaller than the older Torch notebook exports. Shielded training uses the same JAX stage runner as the unshielded baselines; the learned shield internals still load the Torch IOP checkpoint and pickled transition graph that define opponent behavior and environment dynamics.

For shielded IPPO, returned policy adapters are ordinary unmasked base-policy adapters. Runtime safety is still owned by the environment wrapper. During training, however, the wrapper's pre-step mask is part of the behavior policy and therefore part of every stored PPO likelihood.

Cost Handling¶

• IPPO_Lagrangian reads per-agent cost from infos[agent_id]["cost"].
• Missing costs default to 0.0, which matches the labelled env pattern in this repo where only the tracked agent may emit a safety cost.
• Reward and cost GAE are computed separately.
• GAE evaluates every transition's actual next observation. True terminations suppress value bootstrapping; truncations retain it. Both stop the recursive advantage trace, so rollouts containing multiple episodes cannot leak advantages across reset boundaries.
• Reward advantages are standardised over the full rollout and cost advantages are mean-centred by default. The policy loss uses the standard scaled Lagrangian advantage (A_reward - lambda * A_cost) / (1 + lambda).
• Each agent has its own dual variable, initialized by lagrangian_init and clamped to [lambda_min, lambda_max].
• cost_limit is compared with completed, undiscounted episodic cost. Partial episodes persist across rollout boundaries. The dual is updated before the actor update only when at least one episode completed; otherwise it remains unchanged rather than learning from a rollout fragment.
• mean_cost remains a raw per-step diagnostic. cost_return_budget and constraint_violation are the constrained quantities; discounted_cost_return is retained as a critic-aligned diagnostic. dual_update_applied, completed_episodes, and lagrange_multiplier_used expose the exact dual-update lifecycle.
• The single-environment trainers require agents to end an episode together. Asynchronous per-agent termination is rejected explicitly because resetting a shared PettingZoo environment would otherwise corrupt the remaining agents' trajectories.

Notebook Migration¶

• All notebooks/*/*/experiments_mo.py experiment notebooks now route plain IPPO, IPPO-Lagrangian, IPPO-Shielded, and IPPO-TrueShielded training through the JAX stage runner.
• ICPO now uses the repo-native Torch implementation in src/rl/ICPO.py.
• Shielded stages are implemented end to end. Per-environment notebook config still controls whether a given shielded baseline runs by default; some heavyweight or artifact-sensitive runs keep learned run_ippo_shielded=False while exact true shielding remains available.
• src/world_model/sampling/ppo_sampling.py also trains through the JAX IPPO path and exports ordinary Torch nn.Sequential policy and value modules for the Torch-based world/opponent-model pipeline.

Progress Reporting¶

• notebooks/utils/jax_experiments.py owns the notebook-facing training progress bar for run_jax_stage(...).
• The stage reuses a single tqdm.auto.tqdm bar across runs, but resets it to total=timesteps at the start of each run. Notebook cells therefore show the active run's 0..timesteps progress instead of an aggregate num_runs * timesteps counter.
• The description is run-scoped, e.g. IPPO (run 2/3), so the label carries the outer loop while the bar itself reflects the current run's local progress.
• The trainers in src/rl/ippo.py, src/rl/ippo_lag.py, and src/rl/ICPO.py now accept an optional keyword-only progress_callback(completed_steps, total_steps).
• Those trainers report progress once per completed rollout/update cycle, while the notebook wrapper only advances the visible bar at 10% milestones. This keeps long runs observable without the constant notebook refresh churn that per-step updates create.
• The helper also sets a very large tqdm.maxinterval so the background monitor thread does not emit "same step count, later elapsed time" redraws that notebook renderers tend to append as separate lines.
• Set config["progress"] = False in a notebook stage to suppress the bar entirely.

Current boundaries¶

• src/rl/__init__.py now exports only repo-native trainers plus save_models(...).
• src/rl/model_io.py persists the exported Torch modules used by world-model notebooks.
• The old src/rl/old and src/rl/overrides stacks have been deleted. Any new trainer or notebook integration should target src/rl/ICPO.py, src/rl/ippo.py, or src/rl/ippo_lag.py directly.