Update

[ghstack-poisoned]

Author: Vincent Moens

sota-implementations/llm/grpo_mlgym_singlegpu.py

@@ -39,7 +39,7 @@
 parser.add_argument("--batch_size", type=int, default=1)
 parser.add_argument("--epochs", type=int, default=1)
 parser.add_argument("--repeats", type=int, default=16)
-parser.add_argument("--num_envs", type=int, default=32)
+parser.add_argument("--num_envs", type=int, default=1)
 parser.add_argument("--steps_per_batch", type=int, default=64)
 parser.add_argument("--optim_batch_size", type=int, default=4)
 # parser.add_argument("--model_name", type=str, default="gpt2")
@@ -106,6 +106,7 @@ def make_env(tasks: list[str], tokenizer: PreTrainedTokenizer, transform: Transf
         total_dialog_turns=1_000_000,
         async_envs=True,
         replay_buffer=rb,
+        yield_completed_trajectories=True,
     )
     for d in collector:
         print(d)