[Refactor] Fix repeats order

ghstack-source-id: 0bedd5c
Pull Request resolved: #2887

Author: Vincent Moens

test/test_env.py

@@ -4861,31 +4861,22 @@ def policy(td):
         r_reset = r[..., ::max_steps]
         if not batched:
             if str2str:
+                all_strings = r_reset.view(-1)[LLMEnv._DEFAULT_STR_KEY]
+                assert sum(s == all_strings[0] for s in all_strings) == repeats
+                assert sum(s == all_strings[repeats] for s in all_strings) == repeats
                 assert (
-                    r_reset[..., 0][LLMEnv._DEFAULT_STR_KEY]
-                    == r_reset[..., 1][LLMEnv._DEFAULT_STR_KEY]
+                    sum(s == all_strings[repeats * 2] for s in all_strings) == repeats
                 )
+            else:
+                all_tokens = r_reset.view(-1)[LLMEnv._DEFAULT_TOKEN_KEY]
+                assert sum((s == all_tokens[0]).all() for s in all_tokens) == repeats
                 assert (
-                    r_reset[..., 0][LLMEnv._DEFAULT_STR_KEY]
-                    == r_reset[..., 2][LLMEnv._DEFAULT_STR_KEY]
+                    sum((s == all_tokens[repeats]).all() for s in all_tokens) == repeats
                 )
                 assert (
-                    r_reset[..., 0][LLMEnv._DEFAULT_STR_KEY]
-                    != r_reset[..., 3][LLMEnv._DEFAULT_STR_KEY]
+                    sum((s == all_tokens[repeats * 2]).all() for s in all_tokens)
+                    == repeats
                 )
-            else:
-                assert (
-                    r_reset[..., 0][LLMEnv._DEFAULT_TOKEN_KEY]
-                    == r_reset[..., 1][LLMEnv._DEFAULT_TOKEN_KEY]
-                ).all()
-                assert (
-                    r_reset[..., 0][LLMEnv._DEFAULT_TOKEN_KEY]
-                    == r_reset[..., 2][LLMEnv._DEFAULT_TOKEN_KEY]
-                ).all()
-                assert (
-                    r_reset[..., 0][LLMEnv._DEFAULT_TOKEN_KEY]
-                    != r_reset[..., 3][LLMEnv._DEFAULT_TOKEN_KEY]
-                ).any()
         else:
             # When batched, each block contains the 3 reset packs
             if str2str:

torchrl/envs/transforms/llm.py

@@ -562,7 +562,8 @@ def _load_from_dataloader(self, reset: torch.Tensor | None = None):
             if not out.ndim:
                 out = out.unsqueeze(0)
             self._queue.extend(
-                [d for _ in range(max(1, self.repeats)) for d in out.unbind(0)]
+                [d for d in out.unbind(0) for _ in range(max(1, self.repeats))]
+                # [d for _ in range(max(1, self.repeats)) for d in out.unbind(0)]
             )
             return self._queue.popleft()
         return out