Add optional W&B logging to Atari DQN and PPO (#126)

* Add optional W&B logging to Atari PPO

New --wandb CLI flag (in shared parse_args) initializes a wandb run and
logs recent_mean_return, policy_loss, value_loss, and entropy keyed by
global_step. Off by default so non-wandb users see no behavior change.
wandb added to dependencies via uv.

* Gitignore wandb run dir and document --wandb usage

Add wandb/ to .gitignore so local run artifacts don't get committed, and
add a short Logging section to the README covering 'wandb login' and
the --wandb flag. The flag is opt-in and per-user (W&B login is tied to
your own account), so contributors who skip it don't need the package
at runtime.

* Add optional W&B logging to Atari DQN

Same --wandb pattern as PPO: logs recent_mean_return, epsilon, last loss,
and buffer size every 10k frames into project rl-atari-dqn. Off by
default so non-wandb users see no change. README updated to cover both
Atari scripts.

Author: dnddnjs

.gitignore

@@ -7,4 +7,5 @@ __pycache__
 *.egg-info/
 .python-version
 *.pt
+wandb/
 ./Code 2. Cartpole/6. A3C/Cartpole_A3C.pgy

3-atari/1-dqn.py

@@ -117,13 +117,25 @@ def greedy_action(obs):
         online.load_state_dict(torch.load(SAVE_PATH, map_location=device))
         run_test_loop(env, greedy_action)
 
+    if args.wandb:
+        import wandb
+        wandb.init(project="rl-atari-dqn", config={
+            "env": args.env, "total_frames": TOTAL_FRAMES,
+            "buffer_capacity": BUFFER_CAPACITY, "batch_size": BATCH_SIZE,
+            "gamma": GAMMA, "lr": LR, "learn_start": LEARN_START,
+            "train_every": TRAIN_EVERY, "target_update_every": TARGET_UPDATE_EVERY,
+            "epsilon_start": EPSILON_START, "epsilon_end": EPSILON_END,
+            "epsilon_decay_frames": EPSILON_DECAY_FRAMES,
+        })
+
     print(f"device: {device},  env: {args.env},  actions: {n_actions}")
 
     buffer = ReplayBuffer(BUFFER_CAPACITY, env.observation_space.shape)
     obs, _ = env.reset()
     ep_return = 0.0
     recent_returns = deque(maxlen=20)
     train_step = 0
+    last_loss = 0.0
 
     for frame in range(1, TOTAL_FRAMES + 1):
         quit_if_window_closed(env)
@@ -161,16 +173,25 @@ def greedy_action(obs):
             # Gradient clipping — DeepMind uses global norm 10.
             nn.utils.clip_grad_norm_(online.parameters(), 10.0)
             optimizer.step()
+            last_loss = loss.item()
 
             train_step += 1
             if train_step % TARGET_UPDATE_EVERY == 0:
                 target.load_state_dict(online.state_dict())
 
         # Logging.
         if frame % 10_000 == 0:
-            mean = np.mean(recent_returns) if recent_returns else 0.0
+            mean = float(np.mean(recent_returns)) if recent_returns else 0.0
             print(f"frame: {frame:>8}  eps: {epsilon(frame):.3f}  "
                   f"recent_mean_return: {mean:.1f}  buffer: {buffer.size}")
+            if args.wandb:
+                wandb.log({
+                    "global_step": frame,
+                    "recent_mean_return": mean,
+                    "epsilon": epsilon(frame),
+                    "loss": last_loss,
+                    "buffer_size": buffer.size,
+                }, step=frame)
 
     torch.save(online.state_dict(), SAVE_PATH)
     print(f"Saved trained model to {SAVE_PATH}")

3-atari/2-ppo.py

@@ -93,6 +93,16 @@ def policy_action(obs):
     model = ActorCritic(n_actions).to(device)
     optimizer = optim.Adam(model.parameters(), lr=LR, eps=1e-5)
 
+    if args.wandb:
+        import wandb
+        wandb.init(project="rl-atari-ppo", config={
+            "env": args.env, "n_envs": N_ENVS, "rollout_steps": ROLLOUT_STEPS,
+            "total_frames": TOTAL_FRAMES, "epochs": EPOCHS,
+            "minibatch_size": MINIBATCH_SIZE, "clip_coef": CLIP_COEF,
+            "gamma": GAMMA, "gae_lambda": GAE_LAMBDA, "lr": LR,
+            "value_coef": VALUE_COEF, "entropy_coef": ENTROPY_COEF,
+        })
+
     print(f"device: {device},  env: {args.env},  actions: {n_actions},  n_envs: {N_ENVS}")
 
     batch_size = ROLLOUT_STEPS * N_ENVS
@@ -152,6 +162,8 @@ def policy_action(obs):
 
         # --- PPO updates ---
         idx = np.arange(batch_size)
+        pl_sum = vl_sum = ent_sum = 0.0
+        n_mb = 0
         for _ in range(EPOCHS):
             np.random.shuffle(idx)
             for start in range(0, batch_size, MINIBATCH_SIZE):
@@ -173,10 +185,26 @@ def policy_action(obs):
                 nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
                 optimizer.step()
 
+                pl_sum += policy_loss.item()
+                vl_sum += value_loss.item()
+                ent_sum += entropy.item()
+                n_mb += 1
+
+        global_step = update * frames_per_update
         if ep_returns:
             recent = ep_returns[-20:]
-            print(f"update: {update:>4}  frames: {update * frames_per_update:>8}  "
+            print(f"update: {update:>4}  frames: {global_step:>8}  "
                   f"recent_mean_return: {np.mean(recent):.1f}  episodes: {len(ep_returns)}")
+        if args.wandb:
+            log = {
+                "global_step": global_step,
+                "policy_loss": pl_sum / n_mb,
+                "value_loss": vl_sum / n_mb,
+                "entropy": ent_sum / n_mb,
+            }
+            if ep_returns:
+                log["recent_mean_return"] = float(np.mean(ep_returns[-20:]))
+            wandb.log(log, step=global_step)
 
     torch.save(model.state_dict(), SAVE_PATH)
     print(f"Saved trained model to {SAVE_PATH}")

3-atari/env.py

@@ -45,6 +45,8 @@ def parse_args():
                    help="load the saved checkpoint and just play (no learning)")
     p.add_argument("--device", choices=["auto", "cpu", "cuda", "mps"], default="auto",
                    help="override the auto-selected torch device")
+    p.add_argument("--wandb", action="store_true",
+                   help="log metrics to Weights & Biases")
     return p.parse_args()
 
 

README.md

@@ -45,6 +45,18 @@ cd 2-cartpole && uv run python 1-dqn.py --render
 cd 2-cartpole && uv run python 1-dqn.py --test
 ```
 
+### Logging to Weights & Biases (Atari only)
+
+Both Atari scripts (`1-dqn.py`, `2-ppo.py`) can stream training metrics to your own [Weights & Biases](https://wandb.ai/) account. One-time login, then pass `--wandb`:
+
+```bash
+uv run wandb login   # paste the API key from https://wandb.ai/authorize
+cd 3-atari && uv run python 2-ppo.py --env breakout --wandb
+cd 3-atari && uv run python 1-dqn.py --env breakout --wandb
+```
+
+Runs land in *your* `rl-atari-ppo` / `rl-atari-dqn` project — nothing is shared by default. Omit `--wandb` and the script runs without ever touching the network.
+
 ## Updates
 
 Modernized from the 2017 original:

pyproject.toml

@@ -12,4 +12,5 @@ dependencies = [
     "matplotlib>=3.10,<3.11",
     "pygame>=2.6,<3",
     "opencv-python-headless>=4.13,<4.14",
+    "wandb>=0.27.0",
 ]