Bootstrap RSLRL rewards on environment timeouts

motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py

@@ -123,9 +123,9 @@ class RslrlNpEnvWrap(VecEnv):
                          batch_size=[self._num_envs], device=self._device)
 
         # Build extras dict (RSLRL calls it "extras" not "infos")
-        extras = {}
-        if "time_outs" in state.info:
-            extras["time_outs"] = torch.from_numpy(state.info["time_outs"]).to(self._device)
+        extras = {
+            "time_outs": torch.from_numpy(state.truncated.astype(np.float32)).to(self._device),
+        }
 
         # 将 episode 各项奖励传入 TensorBoard（消费后清除，防止重复上报）
         if "ep_report" in state.info:

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