[Feature] TD3 compatibility with compile

ghstack-source-id: fb94307
Pull Request resolved: #2658

Author: Vincent Moens

sota-implementations/cql/cql_online.py

@@ -159,11 +159,7 @@ def update(sampled_tensordict):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     frames_per_batch = cfg.collector.frames_per_batch
     evaluation_interval = cfg.logger.log_interval

sota-implementations/cql/discrete_cql_online.py

@@ -140,11 +140,7 @@ def update(sampled_tensordict):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     eval_rollout_steps = cfg.env.max_episode_steps
     eval_iter = cfg.logger.eval_iter

sota-implementations/crossq/crossq.py

@@ -179,11 +179,7 @@ def update(sampled_tensordict: TensorDict, update_actor: bool):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     eval_iter = cfg.logger.eval_iter
     frames_per_batch = cfg.collector.frames_per_batch

sota-implementations/ddpg/ddpg.py

@@ -145,11 +145,7 @@ def update(sampled_tensordict):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     frames_per_batch = cfg.collector.frames_per_batch
     eval_iter = cfg.logger.eval_iter

sota-implementations/discrete_sac/discrete_sac.py

@@ -144,11 +144,7 @@ def update(sampled_tensordict):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     eval_rollout_steps = cfg.env.max_episode_steps
     eval_iter = cfg.logger.eval_iter

sota-implementations/iql/discrete_iql.py

@@ -148,11 +148,7 @@ def update(sampled_tensordict):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     eval_iter = cfg.logger.eval_iter
     frames_per_batch = cfg.collector.frames_per_batch

sota-implementations/iql/iql_online.py

@@ -145,11 +145,7 @@ def update(sampled_tensordict):
     collected_frames = 0
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     eval_iter = cfg.logger.eval_iter
     frames_per_batch = cfg.collector.frames_per_batch

sota-implementations/sac/config.yaml

@@ -13,7 +13,7 @@ collector:
   frames_per_batch: 1000
   init_env_steps: 1000
   device:
-  env_per_collector: 1
+  env_per_collector: 8
   reset_at_each_iter: False
 
 # replay buffer

sota-implementations/sac/sac.py

@@ -143,11 +143,7 @@ def update(sampled_tensordict):
     pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     init_random_frames = cfg.collector.init_random_frames
-    num_updates = int(
-        cfg.collector.env_per_collector
-        * cfg.collector.frames_per_batch
-        * cfg.optim.utd_ratio
-    )
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
     prb = cfg.replay_buffer.prb
     eval_iter = cfg.logger.eval_iter
     frames_per_batch = cfg.collector.frames_per_batch

sota-implementations/td3/config.yaml

@@ -14,7 +14,7 @@ collector:
   frames_per_batch: 1000
   reset_at_each_iter: False
   device:
-  env_per_collector: 1
+  env_per_collector: 8
   num_workers: 1
 
 # replay buffer
@@ -52,3 +52,8 @@ logger:
   mode: online
   eval_iter: 25000
   video: False
+
+compile:
+  compile: False
+  compile_mode:
+  cudagraphs: False