662af52820
- [x] I have marked all applicable categories:
+ [ ] exception-raising fix
+ [x] algorithm implementation fix
+ [ ] documentation modification
+ [ ] new feature
- [x] I have reformatted the code using `make format` (**required**)
- [x] I have checked the code using `make commit-checks` (**required**)
- [x] If applicable, I have mentioned the relevant/related issue(s)
- [x] If applicable, I have listed every items in this Pull Request
below
While trying to debug Atari PPO+LSTM, I found significant gap between
our Atari PPO example vs [CleanRL's Atari PPO w/
EnvPool](https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_atari_envpoolpy).
I tried to align our implementation with CleaRL's version, mostly in
hyper parameter choices, and got significant gain in Breakout, Qbert,
SpaceInvaders while on par in other games. After this fix, I would
suggest updating our [Atari
Benchmark](https://tianshou.readthedocs.io/en/master/tutorials/benchmark.html)
PPO experiments.
A few interesting findings:
- Layer initialization helps stabilize the training and enable the use
of larger learning rates; without it, larger learning rates will trigger
NaN gradient very quickly;
- ppo.py#L97-L101: this change helps training stability for reasons I do
not understand; also it makes the GPU usage higher.
Shoutout to [CleanRL](https://github.com/vwxyzjn/cleanrl) for a
well-tuned Atari PPO reference implementation!
213 lines
6.5 KiB
Python
213 lines
6.5 KiB
Python
from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Type, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
|
|
from tianshou.utils.net.discrete import NoisyLinear
|
|
|
|
|
|
def layer_init(
|
|
layer: nn.Module, std: float = np.sqrt(2), bias_const: float = 0.0
|
|
) -> nn.Module:
|
|
torch.nn.init.orthogonal_(layer.weight, std)
|
|
torch.nn.init.constant_(layer.bias, bias_const)
|
|
return layer
|
|
|
|
|
|
def scale_obs(module: Type[nn.Module], denom: float = 255.0) -> Type[nn.Module]:
|
|
|
|
class scaled_module(module):
|
|
|
|
def forward(
|
|
self,
|
|
obs: Union[np.ndarray, torch.Tensor],
|
|
state: Optional[Any] = None,
|
|
info: Dict[str, Any] = {}
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
return super().forward(obs / denom, state, info)
|
|
|
|
return scaled_module
|
|
|
|
|
|
class DQN(nn.Module):
|
|
"""Reference: Human-level control through deep reinforcement learning.
|
|
|
|
For advanced usage (how to customize the network), please refer to
|
|
:ref:`build_the_network`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
c: int,
|
|
h: int,
|
|
w: int,
|
|
action_shape: Sequence[int],
|
|
device: Union[str, int, torch.device] = "cpu",
|
|
features_only: bool = False,
|
|
output_dim: Optional[int] = None,
|
|
layer_init: Callable[[nn.Module], nn.Module] = lambda x: x,
|
|
) -> None:
|
|
super().__init__()
|
|
self.device = device
|
|
self.net = nn.Sequential(
|
|
layer_init(nn.Conv2d(c, 32, kernel_size=8, stride=4)),
|
|
nn.ReLU(inplace=True),
|
|
layer_init(nn.Conv2d(32, 64, kernel_size=4, stride=2)),
|
|
nn.ReLU(inplace=True),
|
|
layer_init(nn.Conv2d(64, 64, kernel_size=3, stride=1)),
|
|
nn.ReLU(inplace=True), nn.Flatten()
|
|
)
|
|
with torch.no_grad():
|
|
self.output_dim = np.prod(self.net(torch.zeros(1, c, h, w)).shape[1:])
|
|
if not features_only:
|
|
self.net = nn.Sequential(
|
|
self.net, layer_init(nn.Linear(self.output_dim, 512)),
|
|
nn.ReLU(inplace=True),
|
|
layer_init(nn.Linear(512, np.prod(action_shape)))
|
|
)
|
|
self.output_dim = np.prod(action_shape)
|
|
elif output_dim is not None:
|
|
self.net = nn.Sequential(
|
|
self.net, layer_init(nn.Linear(self.output_dim, output_dim)),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
self.output_dim = output_dim
|
|
|
|
def forward(
|
|
self,
|
|
obs: Union[np.ndarray, torch.Tensor],
|
|
state: Optional[Any] = None,
|
|
info: Dict[str, Any] = {},
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
r"""Mapping: s -> Q(s, \*)."""
|
|
obs = torch.as_tensor(obs, device=self.device, dtype=torch.float32)
|
|
return self.net(obs), state
|
|
|
|
|
|
class C51(DQN):
|
|
"""Reference: A distributional perspective on reinforcement learning.
|
|
|
|
For advanced usage (how to customize the network), please refer to
|
|
:ref:`build_the_network`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
c: int,
|
|
h: int,
|
|
w: int,
|
|
action_shape: Sequence[int],
|
|
num_atoms: int = 51,
|
|
device: Union[str, int, torch.device] = "cpu",
|
|
) -> None:
|
|
self.action_num = np.prod(action_shape)
|
|
super().__init__(c, h, w, [self.action_num * num_atoms], device)
|
|
self.num_atoms = num_atoms
|
|
|
|
def forward(
|
|
self,
|
|
obs: Union[np.ndarray, torch.Tensor],
|
|
state: Optional[Any] = None,
|
|
info: Dict[str, Any] = {},
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
r"""Mapping: x -> Z(x, \*)."""
|
|
obs, state = super().forward(obs)
|
|
obs = obs.view(-1, self.num_atoms).softmax(dim=-1)
|
|
obs = obs.view(-1, self.action_num, self.num_atoms)
|
|
return obs, state
|
|
|
|
|
|
class Rainbow(DQN):
|
|
"""Reference: Rainbow: Combining Improvements in Deep Reinforcement Learning.
|
|
|
|
For advanced usage (how to customize the network), please refer to
|
|
:ref:`build_the_network`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
c: int,
|
|
h: int,
|
|
w: int,
|
|
action_shape: Sequence[int],
|
|
num_atoms: int = 51,
|
|
noisy_std: float = 0.5,
|
|
device: Union[str, int, torch.device] = "cpu",
|
|
is_dueling: bool = True,
|
|
is_noisy: bool = True,
|
|
) -> None:
|
|
super().__init__(c, h, w, action_shape, device, features_only=True)
|
|
self.action_num = np.prod(action_shape)
|
|
self.num_atoms = num_atoms
|
|
|
|
def linear(x, y):
|
|
if is_noisy:
|
|
return NoisyLinear(x, y, noisy_std)
|
|
else:
|
|
return nn.Linear(x, y)
|
|
|
|
self.Q = nn.Sequential(
|
|
linear(self.output_dim, 512), nn.ReLU(inplace=True),
|
|
linear(512, self.action_num * self.num_atoms)
|
|
)
|
|
self._is_dueling = is_dueling
|
|
if self._is_dueling:
|
|
self.V = nn.Sequential(
|
|
linear(self.output_dim, 512), nn.ReLU(inplace=True),
|
|
linear(512, self.num_atoms)
|
|
)
|
|
self.output_dim = self.action_num * self.num_atoms
|
|
|
|
def forward(
|
|
self,
|
|
obs: Union[np.ndarray, torch.Tensor],
|
|
state: Optional[Any] = None,
|
|
info: Dict[str, Any] = {},
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
r"""Mapping: x -> Z(x, \*)."""
|
|
obs, state = super().forward(obs)
|
|
q = self.Q(obs)
|
|
q = q.view(-1, self.action_num, self.num_atoms)
|
|
if self._is_dueling:
|
|
v = self.V(obs)
|
|
v = v.view(-1, 1, self.num_atoms)
|
|
logits = q - q.mean(dim=1, keepdim=True) + v
|
|
else:
|
|
logits = q
|
|
probs = logits.softmax(dim=2)
|
|
return probs, state
|
|
|
|
|
|
class QRDQN(DQN):
|
|
"""Reference: Distributional Reinforcement Learning with Quantile \
|
|
Regression.
|
|
|
|
For advanced usage (how to customize the network), please refer to
|
|
:ref:`build_the_network`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
c: int,
|
|
h: int,
|
|
w: int,
|
|
action_shape: Sequence[int],
|
|
num_quantiles: int = 200,
|
|
device: Union[str, int, torch.device] = "cpu",
|
|
) -> None:
|
|
self.action_num = np.prod(action_shape)
|
|
super().__init__(c, h, w, [self.action_num * num_quantiles], device)
|
|
self.num_quantiles = num_quantiles
|
|
|
|
def forward(
|
|
self,
|
|
obs: Union[np.ndarray, torch.Tensor],
|
|
state: Optional[Any] = None,
|
|
info: Dict[str, Any] = {},
|
|
) -> Tuple[torch.Tensor, Any]:
|
|
r"""Mapping: x -> Z(x, \*)."""
|
|
obs, state = super().forward(obs)
|
|
obs = obs.view(-1, self.action_num, self.num_quantiles)
|
|
return obs, state
|