Upload simoprl/policy.py with huggingface_hub

simoprl/policy.py

@@ -0,0 +1,134 @@
+"""
+Policy network and REINFORCE trainer.
+
+The policy is trained using the *learned* reward model as a surrogate signal.
+Rollouts are generated inside the *real* CartPole environment so that we don't
+compound dynamics-model errors during policy optimisation.
+Evaluation always uses the true environment reward.
+"""
+import torch
+import torch.nn as nn
+import numpy as np
+import gymnasium as gym
+from pathlib import Path
+
+STATE_DIM = 4
+ACTION_DIM = 2
+
+
+class PolicyNetwork(nn.Module):
+    def __init__(self, hidden_dim: int = 64):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(STATE_DIM, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, ACTION_DIM),
+        )
+
+    def forward(self, state_t: torch.Tensor) -> torch.distributions.Categorical:
+        logits = self.net(state_t)
+        return torch.distributions.Categorical(logits=logits)
+
+    def select_action(self, state: np.ndarray) -> tuple[int, torch.Tensor]:
+        s = torch.from_numpy(state.astype(np.float32)).unsqueeze(0)
+        dist = self(s)
+        action = dist.sample()
+        return int(action.item()), dist.log_prob(action)
+
+
+class REINFORCETrainer:
+    """
+    REINFORCE (Williams 1992) using the learned reward model as reward signal.
+
+    Rollouts are collected in the real CartPole environment (not the dynamics
+    model) so we avoid compounding prediction errors. The learned reward model
+    replaces the true reward at training time — the algorithm never sees it.
+    """
+
+    def __init__(
+        self,
+        policy: PolicyNetwork,
+        reward_model,
+        lr: float = 1e-3,
+        gamma: float = 0.99,
+    ):
+        self.policy = policy
+        self.reward_model = reward_model
+        self.optimizer = torch.optim.Adam(policy.parameters(), lr=lr)
+        self.gamma = gamma
+
+    def _collect_episode(self, env) -> tuple[list, list]:
+        """Collect one episode using learned reward, return (log_probs, returns)."""
+        state, _ = env.reset()
+        log_probs, rewards = [], []
+
+        done = False
+        while not done:
+            action, log_prob = self.policy.select_action(state)
+            next_state, _, terminated, truncated, _ = env.step(action)
+            r = self.reward_model.step_reward(state, action)   # learned reward
+            log_probs.append(log_prob)
+            rewards.append(r)
+            state = next_state
+            done = terminated or truncated
+
+        # Discounted returns
+        G, returns = 0.0, []
+        for r in reversed(rewards):
+            G = r + self.gamma * G
+            returns.insert(0, G)
+
+        return log_probs, returns
+
+    def train(self, n_episodes: int = 50) -> None:
+        env = gym.make("CartPole-v1")
+        self.policy.train()
+
+        for ep in range(n_episodes):
+            log_probs, returns = self._collect_episode(env)
+            returns_t = torch.tensor(returns, dtype=torch.float32)
+            returns_t = (returns_t - returns_t.mean()) / (returns_t.std() + 1e-8)
+
+            loss = -torch.stack([lp * R for lp, R in zip(log_probs, returns_t)]).sum()
+            self.optimizer.zero_grad()
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.policy.parameters(), 1.0)
+            self.optimizer.step()
+
+        env.close()
+        self.policy.eval()
+
+    # ── Persistence ───────────────────────────────────────────────────────
+
+    def save(self, path: str) -> None:
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+        torch.save(self.policy.state_dict(), path)
+
+    def load(self, path: str) -> None:
+        self.policy.load_state_dict(torch.load(path, map_location="cpu", weights_only=True))
+
+
+def evaluate_policy(policy: PolicyNetwork, n_episodes: int = 20) -> tuple[float, float]:
+    """
+    Evaluate using the TRUE CartPole environment reward.
+    This is only for measurement — the algorithm never calls this during training.
+    """
+    env = gym.make("CartPole-v1")
+    policy.eval()
+    returns = []
+
+    for _ in range(n_episodes):
+        state, _ = env.reset()
+        total, done = 0.0, False
+        while not done:
+            with torch.no_grad():
+                action, _ = policy.select_action(state)
+            state, reward, terminated, truncated, _ = env.step(action)
+            total += reward
+            done = terminated or truncated
+        returns.append(total)
+
+    env.close()
+    return float(np.mean(returns)), float(np.std(returns))