Correct CPU importance weighting using the IMPALA residual

.github/workflows/install.yml

@@ -99,7 +99,7 @@ jobs:
         run: python tests/test_pufferlib4_final_action_bootstrap.py
 
       - name: Check CPU vector ownership and close behavior
-        run: python -m pytest tests/test_pufferlib4_lifecycle.py tests/test_pufferlib4_layout.py -q
+        run: python -m pytest tests/test_pufferlib4_lifecycle.py tests/test_pufferlib4_layout.py tests/test_pufferlib4_importance.py -q
 
       - name: Check actual PufferLib 4 CPU training and weight reload
         run: python scripts/smoke_pufferlib4_training.py

docs/development/cpu-vtrace.md

@@ -0,0 +1,18 @@
+# CPU importance weighting
+
+Espeholt et al., *IMPALA: Scalable Distributed Deep-RL with Importance Weighted
+Actor-Learner Architectures* (ICML 2018, arXiv:1802.01561), V-trace equations (1)
+and (2), define the clipped importance weight on the entire temporal-difference
+residual: rho * (reward + gamma * next_value - value). The local CPU implementation
+now follows that placement. Existing terminal masking and lambda-weighted trace
+recursion are unchanged.
+
+Primary paper mirror:
+https://www.cl.cam.ac.uk/~ey204/teaching/ACS/R244_2020_2021/papers/Espeholt_ICML_2018.pdf
+
+Tests compare the compiled function with an independent scalar recurrence using
+rho values 0, 0.25, 1 and 2, nonzero values, terminal/nonterminal transitions,
+clipping and a multistep trace. Unit-ratio regression expectations are unchanged.
+The imported CPU build contains no GPU kernels. This addresses the CPU formula
+portion of issue #24; upstream CUDA scalar/SIMD execution and parity remain
+unverified. GPU training must remain unsupported in this package.

src/bindings_cpu.cpp

@@ -61,7 +61,8 @@ static void py_puff_advantage_cpu(
             float r_nxt = rewards[off + t_next];
             float v = values[off + t];
             float v_nxt = values[off + t_next];
-            float delta = rho_t * r_nxt + gamma * v_nxt * nextnonterminal - v;
+            // IMPALA (Espeholt et al., 2018), V-trace TD residual, equation (1).
+            float delta = rho_t * (r_nxt + gamma * v_nxt * nextnonterminal - v);
             lastpufferlam = delta + gamma * lambda * c_t * lastpufferlam * nextnonterminal;
             advantages[off + t] = lastpufferlam;
         }

tests/test_pufferlib4_importance.py

@@ -0,0 +1,24 @@
+"""IMPALA V-trace: rho scales the complete TD residual, not only reward."""
+import numpy as np
+import pytest
+from pufferlib import _C
+
+
+@pytest.mark.parametrize('rho', [0.0, 0.25, 1.0, 2.0])
+@pytest.mark.parametrize('terminal', [0.0, 1.0])
+def test_nonunit_importance_scales_value_residual(rho, terminal):
+    values = np.array([1., 4., 3.], dtype=np.float32)
+    rewards = np.array([0., 6., 2.], dtype=np.float32)
+    dones = np.array([0., terminal, 0.], dtype=np.float32)
+    importance = np.array([rho, 0.5, 1.], dtype=np.float32)
+    advantages = np.zeros(3, dtype=np.float32)
+    _C.puff_advantage_cpu(*(a.ctypes.data for a in [values, rewards, dones, importance, advantages]),
+                          1, 3, 0.5, 0.8, 1., 1.)
+    last = 0.
+    expected = np.zeros(3)
+    for t in reversed(range(2)):
+        nonterminal = 1 - dones[t + 1]
+        residual = min(1., importance[t]) * (rewards[t + 1] + 0.5 * values[t + 1] * nonterminal - values[t])
+        last = residual + 0.5 * 0.8 * min(1., importance[t]) * nonterminal * last
+        expected[t] = last
+    np.testing.assert_allclose(advantages, expected, atol=1e-6)