hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix mock

Browse Source
This commit is contained in:
lucidrains 2025-10-27 10:47:24 -07:00
parent 995b1f64e5
commit 41ab83f691
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
dreamer4/mocks.py
View File

@ -73,8 +73,10 @@ class MockEnv(Module):
out = (state, reward) out = (state, reward)
if self.can_terminate: if self.can_terminate:
shape = (self.num_envs,) if self.vectorized else (1,)
terminate = ( terminate = (
(torch.rand((self.num_envs)) < self.rand_terminate_prob) & (torch.rand(shape) < self.rand_terminate_prob) &
(self._step > self.terminate_after_step) (self._step > self.terminate_after_step)
) )