Skip to content

PR2-IQL Joint-Action Scaling

The May 2026 RESOURCE_EXHAUSTED failure in the six-player gift_refinements/RefinedThenConsume_6_mo PR2_IQL_CER run is explained by the PR2 joint-action critic size, not by replay-buffer storage.

Failing Run

Logs:

  • /homes/oja24/logs/job_238199.out
  • /homes/oja24/logs/job_238251.out

The first run completed and then NashConv was evaluated. The next run failed at the first training update with:

Out of memory while trying to allocate 4.07GiB

The traceback lands at src/rl/pr2_iql.py while converting loss_info JAX arrays to Python floats. That conversion synchronizes the preceding jitted update, so it reports the allocation failure; it is not itself the root cause.

Why The Allocation Is 4.07 GiB

gift_refinements exposes 9 discrete actions per agent. With 6 agents, each PR2-IQL agent has:

  • opponent_joint_size = 9 ** 5 = 59,049
  • critic head width own_action_dim * opponent_joint_size = 9 ** 6 = 531,441

The notebook uses:

  • NUM_STEPS = 256
  • BUFFER_BATCH_SIZE = 4
  • NUM_EPOCHS = 4

Inside src/rl/pr2_iql.py::_make_update_fn, each scanned epoch builds online and target Q tensors for an extended chunk:

(NUM_STEPS + 1) * BUFFER_BATCH_SIZE * 531,441

For this run that is:

257 * 4 * 531,441 * 4 bytes = 2.035 GiB per Q tensor

The online and target tensors together are:

2 * 2.035 GiB = 4.07 GiB

That matches the JAX allocator request exactly. The first run can succeed if the device has enough contiguous memory, but after the run plus NashConv best response compilation/training, JAX/XLA retained allocations and compiled executables make the same 4.07 GiB request fragile.

Chemistry Comparison

The active chemistry notebook was cancelled by Slurm, not by the JAX OOM in the provided traceback. Chemistry has 4 agents and 8 actions, so PR2-IQL's head is:

8 ** 4 = 4,096

With NUM_STEPS = 128 and BUFFER_BATCH_SIZE = 8, online plus target Q tensors are only about 0.031 GiB. Chemistry is slow, but it is not producing the 4.07 GiB PR2-IQL allocation.

Practical Levers

  • Avoid PR2_IQL and PR2_IQL_CER on high-agent-count gridworld substrates unless the joint-action head is redesigned. Its memory grows as action_dim ** num_agents.
  • If a stopgap is needed for gift-refinements, reduce num_steps or buffer_batch_size; both scale this allocation linearly. For example, num_steps=128, buffer_batch_size=2 reduces the online+target Q allocation to about 1.02 GiB.
  • Running one seed per process can avoid some post-NashConv fragmentation and retained executable pressure, but it does not remove the large per-update allocation.
  • TF_GPU_ALLOCATOR=cuda_malloc_async may help fragmentation, as the JAX log suggests, but it is a mitigation rather than a fix.
  • A real PR2-IQL fix would avoid materializing the full [own_action_dim, opponent_joint_action_dim] tensor for every sampled timestep, for example by factorizing opponent response values, sampling opponent actions, or computing only required joint entries plus a reduced bootstrap estimate.