Skip to content

Discrete PR2 Trainers And NashConv Layout

The repo has two discrete-response-policy variants that keep the same top-level trainer contract as the existing PettingZoo trainers:

  • src/rl/pr2_iql.py:make_train(...)
  • src/rl/pr2_ippo.py:make_train(...)

Both run across the repository's discrete-action temporal games. They reuse the notebook-facing hyperparameter blocks that already exist:

  • PR2_IQL reuses the iql block exactly.
  • PR2_IPPO reuses the ippo block exactly.

There are no separate pr2_* notebook hyperparameter sections.

Shared discrete PR2 indexing

src/rl/pr2_common.py owns the discrete opponent joint-action bookkeeping.

build_opponent_indexers(agent_ids, action_dims) builds one OpponentJointActionIndexer per learning agent. Each indexer:

  • keeps the ordered opponent-id tuple for that agent,
  • computes the opponent joint-action size as the product of opponent discrete action counts,
  • encodes an observed opponent action tuple into a single integer index,
  • decodes an integer index back into the opponent action tuple,
  • enumerates opponent tuples in the same order used by the encoder/decoder.

The ordering is general over discrete multi-agent environments. There is no special-cased two-player logic, although the joint-action dimension grows exponentially with the number of opponents.

The important invariant is that replay storage, critic outputs, and explicit enumeration all agree on the same opponent-joint index mapping. The helper now enumerates via decode_opponent_index(index) so the visible tuple order stays aligned with the encoded integer order.

PR2-IQL formulation

PR2_IQL keeps the independent-per-agent structure from IQL, but replaces the scalar action-value head with a joint-action head:

  • one joint-Q network per agent,
  • one target network per agent,
  • one optimizer state per agent,
  • one replay buffer per agent.

For agent i, the critic output is shaped as:

Q_i(s) -> [own_action_dim, opponent_joint_action_dim]

The recursive opponent-response model is derived directly from that critic:

  • rho(a_-i | s, a_i) = softmax(Q_i(s)[a_i, :])
  • q_pr2(s, a_i) = sum_a_-i rho(a_-i | s, a_i) * Q_i(s)[a_i, a_-i]

Action selection is epsilon-greedy over q_pr2(s, a_i), not over raw joint entries. Replay stores:

  • own action,
  • observed opponent joint-action index,
  • reward/done/reset,
  • the final next observation for the sampled chunk.

Updates regress only the observed joint entry:

  • Q_i(s, a_i, a_-i_obs)

The bootstrap target is response-aware:

  • r + gamma * max_a'_i q_pr2_target(s', a'_i) * (1 - done)

All of the outer trainer contract from IQL is preserved:

  • make_train(config, env_factory) -> train(...)
  • warm-start through initial_train_states
  • subset training through trainable_agent_ids
  • update callbacks
  • optional env-history capture
  • the same top-level return keys: train_states, metrics, agent_ids, and optional env_histories

PR2-IPPO formulation

PR2_IPPO keeps the PPO-style rollout/update loop and clipped actor objective, but replaces the scalar state critic with a response-aware joint-Q critic.

For each agent, PR2ActorCritic outputs:

  • actor logits over own actions,
  • joint-Q values shaped as [own_action_dim, opponent_joint_action_dim]

The same critic-derived response model is used:

  • rho(a_-i | s, a_i) = softmax(Q_i(s)[a_i, :])
  • q_pr2(s, a_i) = sum_a_-i rho(a_-i | s, a_i) * Q_i(s)[a_i, a_-i]
  • v_pr2(s) = sum_a_i pi(a_i | s) * q_pr2(s, a_i)

The rollout policy samples from the actor logits. The actor update uses the response-aware advantage from that rollout policy:

  • A_pr2(s, a_taken) = stop_gradient(q_pr2(s, a_taken) - v_pr2(s))

The critic loss regresses the observed joint-action value:

  • Q_i(s, a_taken, a_-i_obs)

toward:

  • r + gamma * v_pr2(s') * (1 - done)

As with IPPO, the trainer contract remains the same at the boundary:

  • warm-starts still restore full train states,
  • frozen-opponent best-response runs still use trainable_agent_ids,
  • metrics remain trainer-owned and shaped like the existing PettingZoo trainer outputs.

The PR2 variants keep temperature internal and fixed at 1.0; no new notebook-facing temperature knob was introduced.

Verifier package split

The old monolithic verifier module has been split into src/verify/nash_conv/.

Current layout:

  • common.py: report helpers, CI helpers, env reset/close helpers
  • iql.py
  • ippo.py
  • nashq.py
  • pr2_iql.py
  • pr2_ippo.py
  • __init__.py: public reexports

Each algorithm-specific module owns:

  • rollout evaluation for that policy family,
  • approximate best-response training using its own trainer family,
  • the unchanged NashConv result payload shape.

src/verify/approx_nash_conv.py remains as a compatibility shim that reexports the public entrypoints from src.verify.nash_conv. Existing imports for the older module path should keep working.

Experiment integration points

The generic temporal-game runner and every current environment adapter treat PR2_IQL and PR2_IPPO as first-class algorithms. Pursuit retains these compatibility helpers:

Public helpers:

  • run_pursuit_pr2_iql_experiment(...)
  • run_pursuit_pr2_ippo_experiment(...)

Notebook suite flags:

  • run_pr2_iql
  • run_pr2_ippo

Algorithm-specific artifact names:

  • save dirs: pr2_iql, pr2_ippo
  • timing keys: pr2_iql_train, pr2_ippo_train
  • comparison labels: PR2_IQL/..., PR2_IPPO/...

Suite configs select these paths with run_pr2_iql and run_pr2_ippo; their CER counterparts use the matching _cer toggles and artifact directories.