[Quality] Longer warmup for cudagraph within sota implementations

ghstack-source-id: 140ba6e
Pull-Request-resolved: #2945

Author: Vincent Moens

sota-implementations/a2c/a2c_atari.py

@@ -183,7 +183,7 @@ def update(batch, max_grad_norm=cfg.optim.max_grad_norm):
         storing_device=device,
         policy_device=device,
         compile_policy={"mode": compile_mode} if cfg.compile.compile else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     # Main loop

sota-implementations/a2c/a2c_mujoco.py

@@ -167,7 +167,7 @@ def update(batch):
         max_frames_per_traj=-1,
         trust_policy=True,
         compile_policy={"mode": compile_mode} if compile_mode is not None else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     test_env.eval()

sota-implementations/discrete_sac/utils.py

@@ -139,7 +139,7 @@ def make_collector(
         device=device,
         storing_device="cpu",
         compile_policy=False if not compile else {"mode": compile_mode},
-        cudagraph_policy=cudagraphs,
+        cudagraph_policy={"warmup": 10} if cudagraphs else False,
     )
     collector.set_seed(cfg.env.seed)
     return collector

sota-implementations/dqn/dqn_atari.py

@@ -178,7 +178,7 @@ def update(sampled_tensordict):
         compile_policy={"mode": compile_mode, "fullgraph": True}
         if compile_mode is not None
         else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     # Main loop

sota-implementations/dqn/dqn_cartpole.py

@@ -136,7 +136,7 @@ def update(sampled_tensordict):
         compile_policy={"mode": compile_mode, "fullgraph": True}
         if compile_mode is not None
         else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     # Main loop

sota-implementations/gail/gail.py

@@ -138,7 +138,7 @@ def main(cfg: DictConfig):  # noqa: F821
         device=device,
         max_frames_per_traj=-1,
         compile_policy={"mode": compile_mode} if compile_mode is not None else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     # Create replay buffer

sota-implementations/iql/utils.py

@@ -138,7 +138,7 @@ def make_collector(cfg, train_env, actor_model_explore, compile_mode):
         total_frames=cfg.collector.total_frames,
         device=device,
         compile_policy={"mode": compile_mode} if compile_mode else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
     collector.set_seed(cfg.env.seed)
     return collector

sota-implementations/ppo/ppo_atari.py

@@ -80,7 +80,7 @@ def main(cfg: DictConfig):  # noqa: F821
         device=device,
         max_frames_per_traj=-1,
         compile_policy={"mode": compile_mode, "warmup": 1} if compile_mode else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     # Create data buffer

sota-implementations/ppo/ppo_mujoco.py

@@ -73,7 +73,7 @@ def main(cfg: DictConfig):  # noqa: F821
         device=device,
         max_frames_per_traj=-1,
         compile_policy={"mode": compile_mode, "warmup": 1} if compile_mode else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
 
     # Create data buffer

sota-implementations/sac/utils.py

@@ -125,7 +125,7 @@ def make_collector(cfg, train_env, actor_model_explore, compile_mode):
         total_frames=cfg.collector.total_frames,
         device=device,
         compile_policy={"mode": compile_mode} if compile_mode else False,
-        cudagraph_policy=cfg.compile.cudagraphs,
+        cudagraph_policy={"warmup": 10} if cfg.compile.cudagraphs else False,
     )
     collector.set_seed(cfg.env.seed)
     return collector