hydra config cartpole

pytorch · vmoens · Dec 7, 2023 · Sep 8, 2023 · Sep 8, 2023 · Sep 8, 2023
commit ece9522208dba32fad3d5f620e385aaa0923ef6b
diff --git a/examples/dqn/config_cartpole.yaml b/examples/dqn/config_cartpole.yaml
@@ -0,0 +1,32 @@
+# Environment
+env:
+  env_name: CartPole-v1
+
+# collector
+collector:
+  total_frames: 500_000
+  frames_per_batch: 10
+  eps_start: 1.0
+  eps_end: 0.05
+  annealing_frames: 250_000
+  init_random_frames: 10_000
+
+# buffer
+buffer:
+  buffer_size: 10_000
+  batch_size: 128
+
+# logger
+logger:
+  backend: csv
+  exp_name: DQN
+
+# Optim
+optim:
+  lr: 2.5e-4
+
+# loss
+loss:
+  gamma: 0.99
+  hard_update_freq: 1
+  num_updates: 1
diff --git a/examples/dqn/dqn_carpole.py b/examples/dqn/dqn_carpole.py
@@ -2,6 +2,7 @@
 DQN Benchmarks: CartPole-v1
 """
 
+import hydra
 import tqdm
 import time
 import torch.nn
@@ -64,81 +65,68 @@ def make_dqn_model(env_name):
     return qvalue_module
 
 
-if __name__ == "__main__":
+@hydra.main(config_path=".", config_name="config_cartpole", version_base="1.1")
+def main(cfg: "DictConfig"):  # noqa: F821
 
     device = "cpu" if not torch.cuda.is_available() else "cuda"
-    env_name = "CartPole-v1"
-    total_frames = 500_000
-    record_interval = 500_000
-    frames_per_batch = 10
-    num_updates = 1
-    buffer_size = 10_000
-    init_random_frames = 10_000
-    annealing_frames = 250_000
-    gamma = 0.99
-    lr = 2.5e-4
-    batch_size = 128
-    hard_update_freq = 50
-    eps_end = 0.05
-    logger_backend = "csv"
-
-    seed = 42
-    torch.manual_seed(seed)
 
     # Make the components
-    model = make_dqn_model(env_name)
-    model_explore = EGreedyWrapper(model, annealing_num_steps=annealing_frames, eps_end=eps_end).to(device)
+    model = make_dqn_model(cfg.env.env_name)
+    model_explore = EGreedyWrapper(
+        policy=model,
+        annealing_num_steps=cfg.collector.annealing_frames,
+        eps_init=cfg.collector.eps_start,
+        eps_end=cfg.collector.eps_end,
+    ).to(device)
 
     # Create the collector
-    collector_class = SyncDataCollector
     collector = SyncDataCollector(
-        make_env(env_name, device),
+        make_env(cfg.env.env_name, device),
         policy=model_explore,
-        frames_per_batch=frames_per_batch,
-        total_frames=total_frames,
+        frames_per_batch=cfg.collector.frames_per_batch,
+        total_frames=cfg.collector.total_frames,
         device=device,
         storing_device=device,
         max_frames_per_traj=-1,
     )
-    collector.set_seed(seed)
 
     # Create the replay buffer
     replay_buffer = TensorDictReplayBuffer(
         pin_memory=False,
         prefetch=3,
         storage=LazyTensorStorage(
-            max_size=buffer_size,
+            max_size=cfg.buffer.buffer_size,
             device=device,
         ),
-        batch_size=batch_size,
+        batch_size=cfg.buffer.batch_size,
     )
 
     # Create the loss module
     loss_module = DQNLoss(
         value_network=model,
-        gamma=gamma,
+        gamma=cfg.loss.gamma,
         loss_function="l2",
         delay_value=True,
     )
-    loss_module.make_value_estimator(gamma=gamma)
-    target_net_updater = HardUpdate(loss_module, value_network_update_interval=hard_update_freq)
+    loss_module.make_value_estimator(gamma=cfg.loss.gamma)
+    target_net_updater = HardUpdate(loss_module, value_network_update_interval=cfg.loss.hard_update_freq)
 
     # Create the optimizer
-    optimizer = torch.optim.Adam(loss_module.parameters(), lr=lr)
+    optimizer = torch.optim.Adam(loss_module.parameters(), lr=cfg.optim.lr)
 
     # Create the logger
-    exp_name = generate_exp_name("DQN", f"CartPole_{env_name}")
-    logger = get_logger(logger_backend, logger_name="dqn", experiment_name=exp_name)
+    exp_name = generate_exp_name("DQN", f"CartPole_{cfg.env.env_name}")
+    logger = get_logger(cfg.logger.backend, logger_name="dqn", experiment_name=exp_name)
 
     # Main loop
     collected_frames = 0
     start_time = time.time()
-    pbar = tqdm.tqdm(total=total_frames)
+    pbar = tqdm.tqdm(total=cfg.collector.total_frames)
 
     for i, data in enumerate(collector):
 
         # Train loging
-        logger.log_scalar("q_values", (data["action_value"]*data["action"]).sum().item() / frames_per_batch, collected_frames)
+        logger.log_scalar("q_values", (data["action_value"]*data["action"]).sum().item() / cfg.collector.frames_per_batch, collected_frames)
         episode_rewards = data["next", "episode_reward"][data["next", "done"]]
         if len(episode_rewards) > 0:
             episode_length = data["next", "step_count"][data["next", "done"]]
@@ -153,10 +141,10 @@ def make_dqn_model(env_name):
         model_explore.step(current_frames)
 
         # optimization steps
-        if collected_frames >= init_random_frames:
-            q_losses = TensorDict({}, batch_size=[num_updates])
-            for j in range(num_updates):
-                sampled_tensordict = replay_buffer.sample(batch_size)
+        if collected_frames >= cfg.collector.init_random_frames:
+            q_losses = TensorDict({}, batch_size=[cfg.collector.num_updates])
+            for j in range(cfg.collector.num_updates):
+                sampled_tensordict = replay_buffer.sample(cfg.buffer.batch_size)
                 loss_td = loss_module(sampled_tensordict)
                 q_loss = loss_td["loss"]
                 optimizer.zero_grad()
@@ -177,3 +165,7 @@ def make_dqn_model(env_name):
     end_time = time.time()
     execution_time = end_time - start_time
     print(f"Training took {execution_time:.2f} seconds to finish")
+
+
+if __name__ == "__main__":
+    main()