Skip to content

IPPO/IQL Cross-Repo Audit

Historical record from the original trainer migration. Treat the current src/rl/ implementations and active algorithm notes as authoritative.

Compared the pltlf PettingZoo IPPO/IQL trainers against the JAX IPPO implementations in /home/omar/Projects/amethyst/csh and /home/omar/Projects/amethyst/omsh.

Findings

  • csh and omsh do not currently contain an IQL implementation under src, tests, or docs; the cross-repo IQL check is therefore limited to pltlf.
  • omsh/src/rl/ippo.py is architecturally closest to a single-env, per-agent trainer class, while pltlf/src/rl/ippo.py keeps a multi-env make_train function. Both use independent actor-critic modules and per-agent optimizers.
  • csh/src/rl/ippo.py uses a homogeneous-agent batched parameter layout and vmap updates. That differs from pltlf's per-agent train-state mapping but is algorithmically still independent PPO.

IQL Replay Contract

pltlf/src/rl/iql.py trains recurrent chunks with this sampled batch schema:

  • obs: [time, batch, obs_dim] current observations
  • reset: [time, batch] recurrent reset masks aligned with obs
  • action, reward, done: [time, batch]
  • final_next_obs: [batch, obs_dim] appended for the final Double-DQN target

The factual replay path must store both obs and final_next_obs. The counterfactual replay path intentionally stores obs_base, obs_state_ids, final_next_base_obs, and final_next_state_id instead, then reconstructs augmented observations after sampling.

Episode Boundaries

The pltlf temporal-logic reward wrappers assign final satisfaction rewards when an episode ends by termination or truncation. IPPO and IQL therefore treat the stored done flag as an episode-boundary mask for return/target computation. This is different from the omsh trainer's explicit bootstrap_dones split, which is useful when a rollout cap is only a collector boundary and should still bootstrap.