Refactor utility functions and update LunarLander version (#26)

* move gae into utils

* move n returns to utils

* move the standard etc into utils

* move gumbel loss into utils

* Rename the lunalancher to v3

* Added devcontainer

* Fix the precommit

* Change to values

* Update the GitHub action Python version to 3.12

Author: Alwaysproblem

.devcontainer/devcontainer.json

@@ -0,0 +1,39 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/anaconda
+{
+  "remoteUser": "root",
+  "name": "tvm",
+  "workspaceMount": "source=${localWorkspaceFolder},target=/root/Desktop/dockerVolumn/${localWorkspaceFolderBasename},type=bind",
+  "workspaceFolder": "/root/Desktop/dockerVolumn/${localWorkspaceFolderBasename}",
+  "image": "alwaysproblem/fastdev-u2204:conda-nv12.2.0",
+  // Features to add to the dev container. More info: https://containers.dev/features.
+  // "features": {},
+  // Use 'forwardPorts' to make a list of ports inside the container available locally.
+  // "forwardPorts": [],
+  // Use 'postCreateCommand' to run commands after the container is created.
+  "postCreateCommand": "bash init.sh",
+  // Configure tool-specific properties.
+  // "customizations": {},
+  // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
+  // "remoteUser": "root"
+  "privileged": true,
+  // "capAdd": ["SYS_PTRACE"],
+  "mounts": [
+    {
+      "source": "/home/lemon/Desktop/dockerVolumn",
+      "target": "/root/Desktop/dockerVolumn",
+      "type": "volume"
+    }
+  ],
+  "runArgs": [
+    // "--cap-add=SYS_PTRACE",
+    // "--security-opt",
+    // "seccomp=unconfined",
+    "--gpus=0",
+    "--shm-size=4G",
+    "memlock=-1:-1",
+    "--ulimit",
+    "--name",
+    "yyx-rltorch",
+  ]
+}

.devcontainer/init.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+setup_new_user 1000 1000
+git config --global --add safe.directory "*"
+
+source /root/miniconda3/etc/profile.d/conda.sh
+
+conda create -n rltorch \
+  pytorch torchvision torchaudio \
+  pytorch-cuda=12.1 gymnasium pyglet \
+  pygame gymnasium-box2d colorama \
+  pylint yapf tqdm 'tensorboardx>=2.5.0' \
+  'tensorboard>2.0' pillow matplotlib scipy \
+  seaborn ipykernel -c conda-forge -c pytorch -c nvidia

.github/workflows/pre-commit.yaml

@@ -12,5 +12,5 @@ jobs:
     - uses: actions/checkout@v3
     - uses: actions/setup-python@v3
       with:
-        python-version: '3.11'
+        python-version: '3.12'
     - uses: pre-commit/action@v3.0.0

.pre-commit-config.yaml

@@ -1,38 +1,26 @@
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.4.0
+  rev: v5.0.0
   hooks:
     - id: check-yaml
     - id: end-of-file-fixer
     - id: trailing-whitespace
     - id: end-of-file-fixer
 
 - repo: https://github.com/pycqa/pylint
-  rev: v2.16.2
+  rev: v3.3.1
   hooks:
     - id: pylint
       args:
         - "--rcfile=.pylintrc"
       exclude: tests(/\w*)*/
 
 - repo: https://github.com/google/yapf
-  rev: v0.40.1
+  rev: v0.40.2
   hooks:
   - id: yapf
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.0.0
+  rev: v1.13.0
   hooks:
     - id: mypy
-
-- repo: https://github.com/pre-commit/mirrors-clang-format
-  rev: 'v15.0.7'
-  hooks:
-  - id: clang-format
-    types_or: [c++, c]
-
-- repo: https://github.com/mwouts/jupytext
-  rev: v1.14.4
-  hooks:
-  - id: jupytext
-    args: [--sync]

.pylintrc

@@ -146,6 +146,7 @@ disable=abstract-method,
         wrong-import-order,
         xrange-builtin,
         zip-builtin-not-iterating,
+        too-many-positional-arguments
 
 
 [REPORTS]

AC/a2c.py

@@ -7,13 +7,7 @@
 
 from util.agent import Agent
 from util.buffer import ReplayBuffer, Trajectory
-
-
-def standardize(v):
-  """Method to standardize a rank-1 np array."""
-  assert len(v) > 1, "Cannot standardize vector of size 1"
-  v_std = (v - v.mean()) / (v.std() + 1e-08)
-  return v_std
+from util.algo import calc_gaes, calc_nstep_return, standardize
 
 
 class Actor(nn.Module):
@@ -211,39 +205,17 @@ def calc_nstep_advs_v_target(self, states, rewards, next_states, terminates):
     with torch.no_grad():
       next_v_pred = self.critic.forward(next_states)
     v_preds = self.critic.forward(states).detach()
-    n_steps_rets = self.calc_nstep_return(
-        rewards=rewards, dones=terminates, next_v_pred=next_v_pred
+    n_steps_rets = calc_nstep_return(
+        rewards=rewards,
+        dones=terminates,
+        next_v_pred=next_v_pred,
+        gamma=self.gamma,
+        n_steps=self.n_steps
     )
     advs = n_steps_rets - v_preds
     v_targets = n_steps_rets
     return standardize(advs), v_targets
 
-  def calc_nstep_return(self, rewards, dones, next_v_pred):
-    T = len(rewards)  #pylint: disable=invalid-name
-    rets = torch.zeros_like(rewards).to(device)
-    _ = 1 - dones
-
-    for i in range(T):
-      # we generate the vector like `gamma = [[γ⁰, γ¹, γ² ...γⁿ]]`
-      # and gamma x reward (vector) to obtain the value for each timestamp.
-      # There are a few items to make it to N
-      # and we will take account all the items.
-      rets[i] = torch.unsqueeze(
-          self.gamma ** torch.arange(len(rewards[i:min(self.n_steps + i, T)])
-                                     ).to(device),
-          dim=0
-      ) @ rewards[i:min(self.n_steps + i, T)]
-
-    if T > self.n_steps:
-      # [[γ⁰, γ¹, γ² ...γⁿ]] x reward.T + γⁿ⁺¹ * V(sₜ₊ₙ₊₁)
-      value_n_steps = self.gamma ** self.n_steps * next_v_pred[self.n_steps:]
-      rets = torch.cat([
-          value_n_steps,
-          torch.zeros(size=(self.n_steps, 1)).to(device)
-      ]) + rets
-
-    return rets
-
   def calc_gae_advs_v_target(self, states, rewards, next_states, terminates):
     """calculate the GAE (Generalized Advantage Estimation) and V_target.
 
@@ -264,23 +236,12 @@ def calc_gae_advs_v_target(self, states, rewards, next_states, terminates):
       next_v_pred = self.critic.forward(next_states[-1])
     v_preds = self.critic.forward(states).detach()
     v_preds_all = torch.concat((v_preds, next_v_pred.unsqueeze(0)), dim=0)
-    advs = self.calc_gaes(rewards, terminates, v_preds_all)
+    advs = calc_gaes(
+        rewards, terminates, v_preds_all, self.gamma, self.gae_lambda
+    )
     v_target = advs + v_preds
     return standardize(advs), v_target
 
-  def calc_gaes(self, rewards, dones, v_preds):
-    # GAE = ∑ₗ (γλ)ˡδₜ₊ₗ
-    # δₜ₊ₗ = rₜ + γV(sₜ₊₁) − V(sₜ)
-    T = len(rewards)  # pylint: disable=invalid-name
-    gaes = torch.zeros_like(rewards, device=device)
-    future_gae = torch.tensor(0.0, dtype=rewards.dtype, device=device)
-    not_dones = 1 - dones  # to reset at episode boundary by multiplying 0
-    deltas = rewards + self.gamma * v_preds[1:] * not_dones - v_preds[:-1]
-    coef = self.gamma * self.gae_lambda
-    for t in reversed(range(T)):
-      gaes[t] = future_gae = deltas[t] + coef * not_dones[t] * future_gae
-    return gaes
-
   def action(self, state, mode="eval"):
     if mode == "train":
       self.actor.train()

AC/main.py

@@ -45,7 +45,7 @@ def main(
   scores = []  # list containing score from each episode
   scores_window = deque(maxlen=100)  # last 100 scores
   eps = eps_start
-  env = gym.make("LunarLander-v2", render_mode="rgb_array")
+  env = gym.make("LunarLander-v3", render_mode="rgb_array")
   # env = gym.make("CartPole-v1", render_mode="rgb_array")
 
   env = TrainMonitor(env, tensorboard_dir="./logs", tensorboard_write_all=True)

AWR/awr.py

@@ -8,23 +8,7 @@
 
 from util.agent import Agent
 from util.buffer import ReplayBuffer, Trajectory
-
-NORMEPS = 1e-8
-
-
-def standardize(v):
-  """Method to standardize a rank-1 np array."""
-  assert len(v) > 1, "Cannot standardize vector of size 1"
-  v_std = (v - v.mean()) / (v.std() + NORMEPS)
-  return v_std
-
-
-def scale_up_values(v, mean=0, std=1, norm_factor=1):
-  return v / norm_factor * std + mean
-
-
-def scale_down_values(v, mean=0, std=1, norm_factor=1):
-  return norm_factor * (v - mean) / (std + NORMEPS)
+from util.algo import standardize, scale_down_values, scale_up_values
 
 
 class Actor(nn.Module):

AWR/main.py

@@ -45,7 +45,7 @@ def main(
   scores = []  # list containing score from each episode
   scores_window = deque(maxlen=100)  # last 100 scores
   eps = eps_start
-  env = gym.make("LunarLander-v2", render_mode="rgb_array")
+  env = gym.make("LunarLander-v3", render_mode="rgb_array")
   # env = gym.make("CartPole-v1", render_mode="rgb_array")
   # max_t = 200
 

DDPG/main.py

@@ -52,7 +52,7 @@ def main(
       render_mode="rgb_array",
   )
   # env = gym.make(
-  #     "LunarLander-v2",
+  #     "LunarLander-v3",
   #     render_mode="rgb_array",
   #     continuous=True,
   # )

DQN/dqn.py

@@ -81,6 +81,8 @@ def __init__(
     self.qnetwork_target = Q(state_dim=state_dims,
                              action_space=action_space).to(device)
     self.optimizer = torch.optim.Adam(self.qnetwork_local.parameters(), lr=lr)
+    # self.optimizer = torch.optim.AdamW(
+    #     self.qnetwork_local.parameters(), lr=lr, amsgrad=True)
 
     # Replay memory
     self.memory = ProportionalPrioritizedReplayBuffer(max_size=mem_size)
@@ -164,7 +166,7 @@ def _learn(self, experiences):
           self.qnetwork_target.forward(next_states).detach(),
           dim=1,
           keepdim=True
-      )[0]
+      ).values
 
     self.memory.update(torch.abs(predicted_targets - labels).squeeze().tolist())
 
@@ -179,6 +181,7 @@ def _learn(self, experiences):
     # loss = self.loss(predicted_targets, labels)
     self.optimizer.zero_grad()
     loss.backward()
+    # torch.nn.utils.clip_grad_value_(self.qnetwork_local.parameters(), 100)
     self.optimizer.step()
 
   def update_targe_q(self):

MPO/main.py

@@ -46,7 +46,7 @@ def main(
   scores_window = deque(maxlen=100)  # last 100 scores
   eps = eps_start
   # env = gym.make("CartPole-v1", render_mode="rgb_array")
-  env = gym.make("LunarLander-v2", render_mode="rgb_array")
+  env = gym.make("LunarLander-v3", render_mode="rgb_array")
 
   env = TrainMonitor(env, tensorboard_dir="./logs", tensorboard_write_all=True)
 

MPO/mpo.py

@@ -9,14 +9,6 @@
 from util.agent import Agent
 from util.buffer import ReplayBuffer, Trajectory
 
-
-def standardize(v):
-  """Method to standardize a rank-1 np array."""
-  assert len(v) > 1, "Cannot standardize vector of size 1"
-  v_std = (v - v.mean()) / (v.std() + 1e-08)
-  return v_std
-
-
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 

PPG/main.py

@@ -46,7 +46,7 @@ def main(
   scores_window = deque(maxlen=100)  # last 100 scores
   eps = eps_start
   # env = gym.make("CartPole-v1", render_mode="rgb_array")
-  env = gym.make("LunarLander-v2", render_mode="rgb_array")
+  env = gym.make("LunarLander-v3", render_mode="rgb_array")
 
   env = TrainMonitor(env, tensorboard_dir="./logs", tensorboard_write_all=True)
 

PPG/ppg.py

@@ -7,13 +7,7 @@
 
 from util.agent import Agent
 from util.buffer import ReplayBuffer, Trajectory, Experience
-
-
-def standardize(v):
-  """Method to standardize a rank-1 np array."""
-  assert len(v) > 1, "Cannot standardize vector of size 1"
-  v_std = (v - v.mean()) / (v.std() + 1e-08)
-  return v_std
+from util.algo import calc_gaes, calc_nstep_return, standardize
 
 
 class Actor(nn.Module):
@@ -367,34 +361,17 @@ def calc_nstep_advs_v_target(self, states, rewards, next_states, terminates):
     with torch.no_grad():
       next_v_pred = self.critic.forward(next_states)
     v_preds = self.critic.forward(states).detach()
-    n_steps_rets = self.calc_nstep_return(
-        rewards=rewards, dones=terminates, next_v_pred=next_v_pred
+    n_steps_rets = calc_nstep_return(
+        rewards=rewards,
+        dones=terminates,
+        next_v_pred=next_v_pred,
+        gamma=self.gamma,
+        n_steps=self.n_steps
     )
     advs = n_steps_rets - v_preds
     v_targets = n_steps_rets
     return standardize(advs), v_targets
 
-  def calc_nstep_return(self, rewards, dones, next_v_pred):
-    T = len(rewards)  #pylint: disable=invalid-name
-    rets = torch.zeros_like(rewards).to(device)
-    _ = 1 - dones
-
-    for i in range(T):
-      rets[i] = torch.unsqueeze(
-          self.gamma ** torch.arange(len(rewards[i:min(self.n_steps + i, T)])
-                                     ).to(device),
-          dim=0
-      ) @ rewards[i:min(self.n_steps + i, T)]
-
-    if T > self.n_steps:
-      value_n_steps = self.gamma ** self.n_steps * next_v_pred[self.n_steps:]
-      rets = torch.cat([
-          value_n_steps,
-          torch.zeros(size=(self.n_steps, 1)).to(device)
-      ]) + rets
-
-    return rets
-
   def calc_gae_advs_v_target(self, states, rewards, next_states, terminates):
     """calculate the GAE (Generalized Advantage Estimation) and V_target.
 
@@ -415,21 +392,12 @@ def calc_gae_advs_v_target(self, states, rewards, next_states, terminates):
       next_v_pred = self.critic.forward(next_states[-1])
     v_preds = self.critic.forward(states).detach()
     v_preds_all = torch.concat((v_preds, next_v_pred.unsqueeze(0)), dim=0)
-    advs = self.calc_gaes(rewards, terminates, v_preds_all)
+    advs = calc_gaes(
+        rewards, terminates, v_preds_all, self.gamma, self.gae_lambda
+    )
     v_target = advs + v_preds
     return standardize(advs), v_target
 
-  def calc_gaes(self, rewards, dones, v_preds):
-    T = len(rewards)  # pylint: disable=invalid-name
-    gaes = torch.zeros_like(rewards, device=device)
-    future_gae = torch.tensor(0.0, dtype=rewards.dtype, device=device)
-    not_dones = 1 - dones  # to reset at episode boundary by multiplying 0
-    deltas = rewards + self.gamma * v_preds[1:] * not_dones - v_preds[:-1]
-    coef = self.gamma * self.gae_lambda
-    for t in reversed(range(T)):
-      gaes[t] = future_gae = deltas[t] + coef * not_dones[t] * future_gae
-    return gaes
-
   def action(self, state, mode="eval"):
     if mode == "train":
       self.actor.train()