Bootstrap RSLRL rewards on environment timeouts

motrix_rl/tests/test_rslrl_np_env_reward_flow.py

@@ -0,0 +1,29 @@
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap
+
+
+def test_step_preserves_rewards_and_marks_only_truncations_as_timeouts():
+    state = SimpleNamespace(
+        obs=np.zeros((2, 3), dtype=np.float32),
+        reward=np.array([1.25, -0.5], dtype=np.float32),
+        done=np.array([True, True]),
+        terminated=np.array([True, False]),
+        truncated=np.array([False, True]),
+        info={},
+    )
+    wrapper = RslrlNpEnvWrap.__new__(RslrlNpEnvWrap)
+    wrapper._env = SimpleNamespace(step=lambda actions: state)
+    wrapper._device = torch.device("cpu")
+    wrapper._state = None
+    wrapper._num_envs = 2
+    wrapper.episode_length_buf = torch.zeros(2, dtype=torch.long)
+
+    _, rewards, dones, extras = wrapper.step(torch.zeros((2, 1)))
+
+    torch.testing.assert_close(rewards, torch.tensor([1.25, -0.5]))
+    torch.testing.assert_close(dones, torch.tensor([1.0, 1.0]))
+    torch.testing.assert_close(extras["time_outs"], torch.tensor([0.0, 1.0]))