REINFORCE

.devcontainer/Dockerfile

@@ -16,7 +16,7 @@ USER $USERNAME
 
 # install poetry for package management
 RUN curl -sSL https://install.python-poetry.org | python3 -
-ENV PATH="~/.local/bin:$PATH"
+ENV PATH="/home/$USERNAME/.local/bin:$PATH"
 
 WORKDIR $DIR 
 

.devcontainer/devcontainer.json

@@ -1,5 +1,5 @@
 {
   "build": { "dockerfile": "Dockerfile", "context": ".." },
   "runArgs": ["--gpus=all"],
-  "extensions": ["ms-python.python", "tamasfe.even-better-toml"],
+  "extensions": ["ms-python.python", "tamasfe.even-better-toml"]
 }

.gitignore

@@ -150,3 +150,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# VSCode
+*.code-workspace

pyproject.toml

@@ -16,12 +16,14 @@ requests = "^2.27.1"
 numpy = "^1.22.0"
 gym = "^0.21.0"
 pandas = "^1.3.5"
+torch = "^1.10.1"
 
 [tool.poetry.dev-dependencies]
 coverage = {extras = ["toml"], version = "^6.2"}
 pytest = "^6.2.5"
 pytest-cov = "^3.0.0"
 pytest-mock = "^3.6.1"
+black = "^21.12b0"
 
 [tool.poetry.scripts]
 

src/functionrl/algorithms/reinforce.py

@@ -0,0 +1,76 @@
+from typing import Optional
+
+import numpy as np
+import torch
+from functionrl.models import LinearNet
+from functionrl.policies import (
+    evaluate_policy,
+    make_categorical_policy_from_model,
+    make_greedy_policy_from_model,
+)
+from torch import optim
+
+from ..envs import make_frozen_lake
+from ..experiences import gen_episodes
+
+
+def reinforce(
+    make_env,
+    gamma: float = 1.0,
+    learning_rate: float = 1e-3,
+    n_episodes: int = 10000,
+    log_interval: int = 100,
+    eval_episodes: int = 1000,
+    seed: Optional[int] = None,
+):
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    env = make_env()
+    n_states = env.observation_space.n
+    n_actions = env.action_space.n
+
+    pi = LinearNet(n_states, n_actions)
+    print(pi)
+
+    optimizer = optim.Adam(pi.parameters(), lr=learning_rate)
+    policy = make_categorical_policy_from_model(pi)
+
+    losses = []
+    for i, episode in enumerate(gen_episodes(env, policy, n=n_episodes), start=1):
+        T = len(episode)
+        rewards = [exp.reward for exp in episode]
+        log_probs = [exp.policy_info["log_prob"] for exp in episode]
+        rets = np.empty(T, dtype=np.float32)
+        future_ret = 0.0
+        for t in reversed(range(T)):
+            future_ret = rewards[t] + gamma * future_ret
+            rets[t] = future_ret
+        rets = torch.tensor(rets)
+        # rets.sub_(rets.mean())
+        log_probs = torch.stack(log_probs)
+        loss = (-log_probs * rets).sum()
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        losses.append(loss.item())
+
+        if i % log_interval == 0:
+            eval_policy = make_greedy_policy_from_model(pi, n_states)
+            mean_return = evaluate_policy(make_env, eval_policy, eval_episodes)
+            mean_loss = np.array(losses[-log_interval:]).mean()
+            print(f"{i:5d}  mean_return: {mean_return:.3f} - loss: {mean_loss:8.4f}")
+
+    return policy
+
+
+if __name__ == "__main__":  # pragma: no cover
+    reinforce(
+        make_frozen_lake,
+        gamma=0.99,
+        learning_rate=0.01,
+        n_episodes=10000,
+        seed=1,
+        eval_episodes=1000,
+    )

src/functionrl/algorithms/tabular_q.py

@@ -1,7 +1,8 @@
+from typing import Optional
 import numpy as np
 from ..utils import linear_decay
-from ..policies import make_epsilon_greedy_policy, make_greedy_policy
-from ..experiences import generate_experiences, generate_episodes
+from ..policies import evaluate_policy, make_epsilon_greedy_policy, make_greedy_policy
+from ..experiences import gen_experiences
 from ..envs import make_frozen_lake
 from ..display import print_pi, print_v
 
@@ -18,48 +19,40 @@ def tabular_q(
     n_steps: int = 5000,
     log_interval: int = 1000,
     eval_episodes: int = 1000,
+    seed: Optional[int] = None,
 ):
-    env_train = make_env()
-    env_eval = make_env()
+    env = make_env()
 
-    n_states = env_train.observation_space.n
-    n_actions = env_train.action_space.n
+    n_states = env.observation_space.n
+    n_actions = env.action_space.n
+
+    q = np.zeros((n_states, n_actions))
 
     alpha_decay = linear_decay(alpha_max, alpha_min, alpha_decay_steps)
     epsilon_decay = linear_decay(epsilon_max, epsilon_min, epsilon_decay_steps)
 
-    q = np.zeros((n_states, n_actions))
-
-    # TODO: pass decay into make_eps
-    policy_train = make_epsilon_greedy_policy(
-        q, epsilon_max, epsilon_min, epsilon_decay_steps
-    )
+    policy_train = make_epsilon_greedy_policy(q, epsilon_decay, seed=seed)
     policy_eval = make_greedy_policy(q)
 
-    for step, exp in enumerate(
-        generate_experiences(env_train, policy_train, n=n_steps)
-    ):
-
-        td_target = (
-            exp.reward + gamma * float(not exp.is_done) * q[exp.next_state].max()
-        )
-        td_error = td_target - q[exp.state, exp.action]
+    for i, exp in enumerate(gen_experiences(env, policy_train, n=n_steps), start=1):
+        state, action, reward, next_state, is_done, policy_info = exp
+        td_target = reward + gamma * float(not is_done) * q[next_state].max()
+        td_error = td_target - q[state, action]
 
-        alpha = alpha_decay(step)
-        q[exp.state, exp.action] += alpha * td_error
+        alpha = alpha_decay(i)
+        q[state, action] += alpha * td_error
 
-        if (step + 1) % log_interval == 0:
-            episodes = list(generate_episodes(env_eval, policy_eval, n=eval_episodes))
-            returns = [sum(e.reward for e in episode) for episode in episodes]
-            mean_return = np.mean(returns)
-            print(
-                f"{step+1:5d}: {mean_return:.3f}, eps: {epsilon_decay(step):.3f}, alpha: {alpha:.6f}"
-            )
+        if i % log_interval == 0:
+            epsilon = policy_info["epsilon"]
+            mean_return = evaluate_policy(make_env, policy_eval, eval_episodes)
+            print(f"{i:5d}: {mean_return:.3f}, eps: {epsilon:.3f}, alpha: {alpha:.6f}")
+            pi = np.argmax(q, axis=1)
+            print_pi(pi)
 
     return q
 
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     q = tabular_q(
         make_frozen_lake,
         gamma=1,
@@ -71,8 +64,7 @@ def tabular_q(
         epsilon_decay_steps=100_000,
         n_steps=100_000,
         log_interval=10_000,
+        seed=0,
     )
-    pi = np.argmax(q, axis=1)
-    print_pi(pi)
     v = np.max(q, axis=1)
     print_v(v)