Skip to content

Shared RL Infrastructure

The PettingZoo trainers share environment/runtime contracts in src/rl/common.py and trajectory/replay contracts in src/rl/trajectory.py. Algorithm modules should import these public helpers instead of importing private names from another trainer.

rl/common.py

The common module owns:

  • EnvFactory and update-callback types;
  • deterministic integer/JAX-key seed coercion;
  • strict modern ParallelEnv.reset() handling;
  • observation flattening and Box/Discrete space validation;
  • agent-subset validation and Flax train-state restoration;
  • environment history extraction, closing, and optional W&B logging; and
  • reusable numeric and rollout-budget configuration validators.

The reset contract intentionally requires (observations, infos). Legacy single-value reset results are rejected rather than guessed. Whole-environment autoreset remains the trainer-level convention when any agent terminates or truncates.

rl/trajectory.py

The trajectory module owns categorical log-probability/entropy, masked means and normalization, finite-episode GAE, clipped value loss, and contiguous sequence replay.

SequenceReplayBuffer has these invariants:

  • capacity is counted in transitions;
  • the required field schema is exact and fixed after first insertion;
  • trailing shapes and dtypes cannot change;
  • every insertion declares a logical stream;
  • per-stream step numbers establish temporal adjacency independently of the physical ring-buffer position;
  • sampling never crosses a stream switch or overwritten step-number gap;
  • missing burn-in context is legal only when an explicit reset establishes the recurrent boundary; and
  • returned loss_mask distinguishes burn-in/padding from learning steps.

Replay samples use replacement. sequence_valid describes real rather than left-padded positions and is available for diagnostics; learner loss is driven by loss_mask.

Configuration ownership

validate_ippo_config() and validate_iql_config() live beside their algorithm families because they also establish derived algorithm fields. PR2-IPPO reuses the IPPO validator and PR2-IQL reuses the IQL validator. The temporal-game resolver maps notebook replay_sequence_length and burn_in_steps fields to the uppercase trainer config.

Timestep budgets must be exact multiples of the rollout batch. Silent flooring would otherwise train for fewer transitions than requested and desynchronize schedule/resume accounting.

Alternatives considered

Storing rollout-start hidden states was rejected because replayed hidden states would have been produced by stale parameters. Recomputing from a contiguous burn-in prefix uses the current online and target parameters instead. Sampling arbitrary transitions was retained only for feedforward behavior through a one-step or longer learning window with zero burn-in; recurrent learners always use stream-aware histories.