Skip to content

PettingZoo IQL Trainer

src/rl/iql.py now targets PettingZoo ParallelEnv instances directly, in the same broad style as src/rl/ippo.py.

Supported entrypoint

Use:

train = make_train(config, env_factory)
out = train(seed=0)
  • config is the flat IQL hyperparameter dictionary.
  • env_factory(seed) must return a fresh PettingZoo ParallelEnv.

The returned trainer callable also accepts:

out = train(
    seed=0,
    initial_train_states=...,
    trainable_agent_ids={"player_0"},
    on_update=callback,
    capture_env_histories=True,
)
  • initial_train_states: warm-start all agents from a prior IQL snapshot.
  • trainable_agent_ids: train only that subset; the rest act as frozen greedy opponents.
  • on_update(update_idx, total_updates, metric): called once per IQL update.
  • capture_env_histories=True: serializes wrapper histories into env_histories.

Core design

  • One independent Q-network per agent.
  • One independent target network per agent.
  • One optimizer state per agent.
  • One replay buffer per agent.
  • No parameter sharing.

The trainer currently assumes:

  • PettingZoo ParallelEnv
  • gymnasium.spaces.Box observations
  • gymnasium.spaces.Discrete actions
  • whole-environment autoreset on any agent termination or truncation

That matches the current experiment environments and monitor wrappers.

Network kinds

Set config["NETWORK_KIND"] to:

  • "rnn": default; keeps the GRU-style recurrent Q-network shape from the older IQL code.
  • "mlp": feedforward Q-network with the same trainer contract.

Both modes still keep one hidden-state slot per env and per agent so the action selection path can stay structurally uniform.

Contiguous transition replay

IQL and PR2-IQL use SequenceReplayBuffer from src/rl/trajectory.py. BUFFER_SIZE is an exact transition capacity, not a chunk count. Every stored transition has the fixed factual schema obs, reset, action, reward, done, and next_obs; PR2-IQL adds opponent_index.

Every insertion has a logical stream id. Factual transitions retain one stream per environment across rollout updates. Sampling sorts by per-stream sequence number and rejects windows that cross a stream boundary or a ring-buffer gap. This is what makes recurrent history valid even when a learning sequence spans several rollout collections.

CER variants are materialized into the same fixed schema, but every synthetic variant gets a fresh stream and a forced reset at its first transition. The factual CER variant continues the real environment stream. This prevents two unrelated alternative monitor histories from being concatenated into a fake recurrent history.

Recurrent burn-in and truncated BPTT

The learning-window length is REPLAY_SEQUENCE_LENGTH; it defaults to NUM_STEPS. BURN_IN_STEPS defaults to min(16, REPLAY_SEQUENCE_LENGTH) for recurrent networks and must be zero for MLPs.

Each sample contains up to BURN_IN_STEPS real transitions immediately before the learning window. At a known episode start, missing left context is padded with reset-marked zero observations. A window with missing history and no known reset is ineligible. Replay therefore never silently treats the middle of an episode as an initial recurrent state.

The update:

  1. initializes the online and target GRU states to zero;
  2. unrolls both networks over the burn-in prefix with their current parameters;
  3. applies stop_gradient to both reconstructed hidden states;
  4. unrolls the learning observations plus the final next_obs; and
  5. applies the TD loss only where the learning-window mask is one.

This is truncated backpropagation through time: burn-in reconstructs the current hidden state without propagating gradients through older transitions. The base IQL target is Double DQN: the online network selects the next action and the target network evaluates it. PR2-IQL applies the same burn-in contract to its response-aware joint-action target.

The replay capacity validator requires room for one learning-plus-burn-in window per environment. The initial observation after env.reset() and every autoreset observation carry reset=1.

Configuration validation

validate_iql_config() is shared by IQL and PR2-IQL. It rejects missing required fields, non-integral counts, invalid ranges, incompatible MLP burn-in, undersized replay, reversed epsilon schedules, and timestep budgets that are not exactly divisible by NUM_ENVS * NUM_STEPS. It returns the normalized derived fields used by the trainer.

Runtime shape

The IQL and PR2-IQL trainers keep the Python environment loop, but avoid tiny JAX calls inside the hot path:

  • action selection is compiled per agent and handles the one-step recurrent forward pass, epsilon-greedy sampling, and greedy evaluation actions;
  • sequence replay samples for all NUM_EPOCHS are stacked with a leading epoch axis;
  • the update JIT scans over that epoch axis, so each trainable agent pays one compiled update call per environment update rather than one call per epoch.

For the pursuit notebooks this keeps the visible progress unit unchanged: NUM_ENVS * NUM_STEPS still equals the env-step increment reported to tqdm. The change is only in how much host/device overhead IQL pays between progress bar redraws.

Returned outputs

The trainer returns an IPPO-compatible shape:

  • train_states
  • metrics
  • agent_ids
  • optional env_histories

Shared metric keys include:

  • episode_return_mean
  • episode_length_mean
  • completed_episodes
  • per_agent_episode_return_mean

IQL also reports:

  • loss
  • qvals
  • epsilon

IQL NashConv BR settings

approximate_nashconv_bernoulli_iql(...) uses IQL itself to search for better best responses.

For each deviating agent it rebuilds a fresh IQL trainer with:

  • NUM_ENVS = 1
  • NUM_STEPS = br_batch_steps
  • TOTAL_TIMESTEPS = br_train_iters * br_batch_steps
  • BUFFER_BATCH_SIZE = 1
  • LEARNING_STARTS = 0
  • TEST_DURING_TRAINING = False

All non-target agents are warm-started from the snapshot and kept frozen via trainable_agent_ids={target_agent}.