From 57faf8880b690fb700d576581d44bb216d116f50 Mon Sep 17 00:00:00 2001
From: 8x54zj-m <8x54zj-m@motrixlab.local>
Date: Wed, 22 Jul 2026 13:30:25 +0800
Subject: [PATCH] Bootstrap RSLRL rewards on environment timeouts

---
 .../src/motrix_rl/rslrl/torch/wrap_vec_env.py |  6 ++--
 .../tests/test_rslrl_np_env_reward_flow.py    | 29 +++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)
 create mode 100644 motrix_rl/tests/test_rslrl_np_env_reward_flow.py

diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py b/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py
index 5c7ebb7..fd3dab8 100644
--- a/motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py
+++ b/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:
diff --git a/motrix_rl/tests/test_rslrl_np_env_reward_flow.py b/motrix_rl/tests/test_rslrl_np_env_reward_flow.py
new file mode 100644
index 0000000..88a07cb
--- /dev/null
+++ b/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]))