Skip to content

PettingZoo Nash-Q Trainer

src/rl/nash_q.py adds a tabular Nash-Q trainer that follows the same high-level contract as the PettingZoo IPPO and IQL trainers:

train = make_train(config, env_factory)
out = train(seed=0)

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,
)

Config shape

The flat Nash-Q config keys are:

  • NUM_ENVS
  • TOTAL_TIMESTEPS
  • NUM_STEPS
  • GAMMA
  • EPS_START
  • EPS_FINISH
  • EPS_DECAY
  • MIN_ALPHA
  • EQUILIBRIUM_SELECTION
  • SEED

NUM_UPDATES is derived internally as TOTAL_TIMESTEPS // NUM_STEPS // NUM_ENVS.

State and value representation

  • The trainer assumes a fixed possible_agents order from the PettingZoo env.
  • Each state key is the ordered tuple of per-agent observation byte strings.
  • Each agent keeps one tabular Q_i(s, a_0, ..., a_{n-1}).
  • Joint-action visit counts are shared across agents for the harmonic step size alpha = max(1 / N(s, a), MIN_ALPHA).

The implementation requires Discrete action spaces. Observations only need to round-trip through np.asarray(..., dtype=np.float32), which covers the matrix and Pursuit environments and the flat object vectors used by native gridworld experiments.

Equilibrium solver path

Stage-game equilibria are solved from the per-agent payoff tensors with pygambit, which is a project dependency:

  1. pygambit.nash.enumpure_solve(...)
  2. pygambit.nash.enumpoly_solve(..., use_strategic=True, stop_after=1)

There is intentionally no uniform product-policy fallback. If pygambit cannot build or solve a state game, Nash-Q raises instead of silently changing the backup target.

EQUILIBRIUM_SELECTION controls tie-breaking among equilibria returned by a solver:

  • payoff_dominant (default): choose the returned equilibrium with the largest total expected Q value, then the largest minimum per-agent expected Q value, then the earliest solver-returned equilibrium.
  • first: keep the older behavior and use the first solver-returned equilibrium.

This matters for weakly coordinated games such as Matrix Stag Hunt. With the older first-pure rule, Hare/Hare is selected before Stag/Stag even after the learned Stag/Stag value becomes higher, because Hare/Hare remains a weak Nash equilibrium. The payoff-dominant default lets positive coordinated values propagate backward through Nash-Q backups.

Snapshot format

Nash-Q snapshots are plain Python mappings rather than Flax TrainStates. Each agent entry stores:

  • q_values: dict[state_key, np.ndarray]
  • visit_counts: dict[state_key, np.ndarray]
  • timesteps
  • n_updates

The generic experiment checkpoint bundle preserves those mappings directly so tuple state keys and NumPy arrays survive pickling unchanged.

Best-response mode

When trainable_agent_ids contains exactly one agent, the trainer switches to a best-response style update:

  • frozen opponents act from the equilibrium policies induced by the frozen baseline snapshot,
  • the trainable agent acts from a tie-broken mixed best response to those opponent policies,
  • the Bellman target uses the best-response continuation value against the same frozen opponent mix.

approximate_nashconv_bernoulli_nashq(...) reuses that path by warm-starting from the baseline snapshot and training only the deviating agent.

Returned outputs

The trainer returns:

  • train_states
  • metrics
  • agent_ids
  • optional env_histories

Shared metric keys:

  • episode_return_mean
  • episode_length_mean
  • completed_episodes
  • per_agent_episode_return_mean

Nash-Q-specific metrics:

  • epsilon