config

pytorch · Sep 19, 2023 · 663acbd · 663acbd
1 parent 888fbcb
commit 663acbd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/examples/impala/config.yaml b/examples/impala/config.yaml
@@ -4,8 +4,9 @@ env:
 
 # collector
 collector:
-  frames_per_batch: 80
+  frames_per_batch: 128 # 80
   total_frames: 40_000_000
+  num_workers: 8
 
 # logger
 logger:
@@ -16,8 +17,8 @@ logger:
 
 # Optim
 optim:
-  lr: 0.0001
-  eps: 1.0e-8
+  lr: 0.0006 # 0.0001
+  eps: 1.0e-5
   weight_decay: 0.0
   max_grad_norm: 40.0
   anneal_lr: True

diff --git a/examples/impala/impala.py b/examples/impala/impala.py
@@ -25,7 +25,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
     from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
     from torchrl.envs import ExplorationType, set_exploration_type
-    from torchrl.objectives import A2CLoss, ClipPPOLoss
+    from torchrl.objectives import A2CLoss
     from torchrl.record.loggers import generate_exp_name, get_logger
     from torchrl.objectives.value.vtrace import VTrace
     from utils import make_parallel_env, make_ppo_models
@@ -58,7 +58,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     #     sync=False,
     # )
     collector = MultiaSyncDataCollector(
-        create_env_fn=[make_parallel_env(cfg.env.env_name, device)] * 8,
+        create_env_fn=[make_parallel_env(cfg.env.env_name, device)] * cfg.collector.num_workers,
         policy=actor,
         frames_per_batch=frames_per_batch,
         total_frames=total_frames,