fix mock

2025-10-27 10:47:24 -07:00 · 2025-10-27 10:47:24 -07:00 · 41ab83f691
commit 41ab83f691
parent 995b1f64e5
1 changed files with 3 additions and 1 deletions
--- a/dreamer4/mocks.py
+++ b/dreamer4/mocks.py
@ -73,8 +73,10 @@ class MockEnv(Module):
        out = (state, reward)

        if self.can_terminate:
+            shape = (self.num_envs,) if self.vectorized else (1,)
+
            terminate = (
-                (torch.rand((self.num_envs)) < self.rand_terminate_prob) &
+                (torch.rand(shape) < self.rand_terminate_prob) &
                (self._step > self.terminate_after_step)
            )