pushing model
Browse files- .gitattributes +1 -0
- README.md +93 -0
- events.out.tfevents.1675616243.ip-26-0-142-102.1462862.0 +3 -0
- poetry.lock +0 -0
- pyproject.toml +178 -0
- replay.mp4 +0 -0
- sebulba_ppo_envpool.cleanrl_model +3 -0
- sebulba_ppo_envpool.py +889 -0
- videos/Pitfall-v5__sebulba_ppo_envpool__1__552a94c8-c755-4d9a-a517-2968c95f4ca4-eval/0.mp4 +0 -0
.gitattributes
CHANGED
@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
32 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
33 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
34 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
32 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
33 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
34 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
35 |
+
sebulba_ppo_envpool.cleanrl_model filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
tags:
|
3 |
+
- Pitfall-v5
|
4 |
+
- deep-reinforcement-learning
|
5 |
+
- reinforcement-learning
|
6 |
+
- custom-implementation
|
7 |
+
library_name: cleanrl
|
8 |
+
model-index:
|
9 |
+
- name: PPO
|
10 |
+
results:
|
11 |
+
- task:
|
12 |
+
type: reinforcement-learning
|
13 |
+
name: reinforcement-learning
|
14 |
+
dataset:
|
15 |
+
name: Pitfall-v5
|
16 |
+
type: Pitfall-v5
|
17 |
+
metrics:
|
18 |
+
- type: mean_reward
|
19 |
+
value: 0.00 +/- 0.00
|
20 |
+
name: mean_reward
|
21 |
+
verified: false
|
22 |
+
---
|
23 |
+
|
24 |
+
# (CleanRL) **PPO** Agent Playing **Pitfall-v5**
|
25 |
+
|
26 |
+
This is a trained model of a PPO agent playing Pitfall-v5.
|
27 |
+
The model was trained by using [CleanRL](https://github.com/vwxyzjn/cleanrl) and the most up-to-date training code can be
|
28 |
+
found [here](https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/sebulba_ppo_envpool.py).
|
29 |
+
|
30 |
+
## Get Started
|
31 |
+
|
32 |
+
To use this model, please install the `cleanrl` package with the following command:
|
33 |
+
|
34 |
+
```
|
35 |
+
pip install "cleanrl[jax,envpool,atari]"
|
36 |
+
python -m cleanrl_utils.enjoy --exp-name sebulba_ppo_envpool --env-id Pitfall-v5
|
37 |
+
```
|
38 |
+
|
39 |
+
Please refer to the [documentation](https://docs.cleanrl.dev/get-started/zoo/) for more detail.
|
40 |
+
|
41 |
+
|
42 |
+
## Command to reproduce the training
|
43 |
+
|
44 |
+
```bash
|
45 |
+
curl -OL https://huggingface.co/cleanrl/Pitfall-v5-sebulba_ppo_envpool-seed1/raw/main/sebulba_ppo_envpool.py
|
46 |
+
curl -OL https://huggingface.co/cleanrl/Pitfall-v5-sebulba_ppo_envpool-seed1/raw/main/pyproject.toml
|
47 |
+
curl -OL https://huggingface.co/cleanrl/Pitfall-v5-sebulba_ppo_envpool-seed1/raw/main/poetry.lock
|
48 |
+
poetry install --all-extras
|
49 |
+
python sebulba_ppo_envpool.py --actor-device-ids 0 --learner-device-ids 1 2 3 4 --params-queue-timeout 0.02 --track --save-model --upload-model --hf-entity cleanrl --env-id Pitfall-v5 --seed 1
|
50 |
+
```
|
51 |
+
|
52 |
+
# Hyperparameters
|
53 |
+
```python
|
54 |
+
{'actor_device_ids': [0],
|
55 |
+
'anneal_lr': True,
|
56 |
+
'async_batch_size': 16,
|
57 |
+
'async_update': 4,
|
58 |
+
'batch_size': 8192,
|
59 |
+
'capture_video': False,
|
60 |
+
'clip_coef': 0.1,
|
61 |
+
'cuda': True,
|
62 |
+
'ent_coef': 0.01,
|
63 |
+
'env_id': 'Pitfall-v5',
|
64 |
+
'exp_name': 'sebulba_ppo_envpool',
|
65 |
+
'gae_lambda': 0.95,
|
66 |
+
'gamma': 0.99,
|
67 |
+
'hf_entity': 'cleanrl',
|
68 |
+
'learner_device_ids': [1, 2, 3, 4],
|
69 |
+
'learning_rate': 0.00025,
|
70 |
+
'max_grad_norm': 0.5,
|
71 |
+
'minibatch_size': 2048,
|
72 |
+
'norm_adv': True,
|
73 |
+
'num_actor_threads': 1,
|
74 |
+
'num_envs': 64,
|
75 |
+
'num_minibatches': 4,
|
76 |
+
'num_steps': 128,
|
77 |
+
'num_updates': 6103,
|
78 |
+
'params_queue_timeout': 0.02,
|
79 |
+
'profile': False,
|
80 |
+
'save_model': True,
|
81 |
+
'seed': 1,
|
82 |
+
'target_kl': None,
|
83 |
+
'test_actor_learner_throughput': False,
|
84 |
+
'torch_deterministic': True,
|
85 |
+
'total_timesteps': 50000000,
|
86 |
+
'track': True,
|
87 |
+
'update_epochs': 4,
|
88 |
+
'upload_model': True,
|
89 |
+
'vf_coef': 0.5,
|
90 |
+
'wandb_entity': None,
|
91 |
+
'wandb_project_name': 'cleanRL'}
|
92 |
+
```
|
93 |
+
|
events.out.tfevents.1675616243.ip-26-0-142-102.1462862.0
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:3300785340e869e22364c52de5d655c09eba9a97806fad8c2dfd570087459ad4
|
3 |
+
size 9358614
|
poetry.lock
ADDED
The diff for this file is too large to render.
See raw diff
|
|
pyproject.toml
ADDED
@@ -0,0 +1,178 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
[tool.poetry]
|
2 |
+
name = "cleanrl"
|
3 |
+
version = "1.1.0"
|
4 |
+
description = "High-quality single file implementation of Deep Reinforcement Learning algorithms with research-friendly features"
|
5 |
+
authors = ["Costa Huang <costa.huang@outlook.com>"]
|
6 |
+
packages = [
|
7 |
+
{ include = "cleanrl" },
|
8 |
+
{ include = "cleanrl_utils" },
|
9 |
+
]
|
10 |
+
keywords = ["reinforcement", "machine", "learning", "research"]
|
11 |
+
license="MIT"
|
12 |
+
readme = "README.md"
|
13 |
+
|
14 |
+
[tool.poetry.dependencies]
|
15 |
+
python = ">=3.7.1,<3.10"
|
16 |
+
tensorboard = "^2.10.0"
|
17 |
+
wandb = "^0.13.6"
|
18 |
+
gym = "0.23.1"
|
19 |
+
torch = ">=1.12.1"
|
20 |
+
stable-baselines3 = "1.2.0"
|
21 |
+
gymnasium = "^0.26.3"
|
22 |
+
moviepy = "^1.0.3"
|
23 |
+
pygame = "2.1.0"
|
24 |
+
huggingface-hub = "^0.11.1"
|
25 |
+
|
26 |
+
ale-py = {version = "0.7.4", optional = true}
|
27 |
+
AutoROM = {extras = ["accept-rom-license"], version = "^0.4.2"}
|
28 |
+
opencv-python = {version = "^4.6.0.66", optional = true}
|
29 |
+
pybullet = {version = "3.1.8", optional = true}
|
30 |
+
procgen = {version = "^0.10.7", optional = true}
|
31 |
+
pytest = {version = "^7.1.3", optional = true}
|
32 |
+
mujoco = {version = "^2.2", optional = true}
|
33 |
+
imageio = {version = "^2.14.1", optional = true}
|
34 |
+
free-mujoco-py = {version = "^2.1.6", optional = true}
|
35 |
+
mkdocs-material = {version = "^8.4.3", optional = true}
|
36 |
+
markdown-include = {version = "^0.7.0", optional = true}
|
37 |
+
jax = {version = "^0.3.17", optional = true}
|
38 |
+
jaxlib = {version = "^0.3.15", optional = true}
|
39 |
+
flax = {version = "^0.6.0", optional = true}
|
40 |
+
optuna = {version = "^3.0.1", optional = true}
|
41 |
+
optuna-dashboard = {version = "^0.7.2", optional = true}
|
42 |
+
rich = {version = "<12.0", optional = true}
|
43 |
+
envpool = {version = "^0.8.1", optional = true}
|
44 |
+
PettingZoo = {version = "1.18.1", optional = true}
|
45 |
+
SuperSuit = {version = "3.4.0", optional = true}
|
46 |
+
multi-agent-ale-py = {version = "0.1.11", optional = true}
|
47 |
+
boto3 = {version = "^1.24.70", optional = true}
|
48 |
+
awscli = {version = "^1.25.71", optional = true}
|
49 |
+
shimmy = {version = "^0.1.0", optional = true}
|
50 |
+
dm-control = {version = "^1.0.8", optional = true}
|
51 |
+
|
52 |
+
[tool.poetry.group.dev.dependencies]
|
53 |
+
pre-commit = "^2.20.0"
|
54 |
+
|
55 |
+
[tool.poetry.group.atari]
|
56 |
+
optional = true
|
57 |
+
[tool.poetry.group.atari.dependencies]
|
58 |
+
ale-py = "0.7.4"
|
59 |
+
AutoROM = {extras = ["accept-rom-license"], version = "^0.4.2"}
|
60 |
+
opencv-python = "^4.6.0.66"
|
61 |
+
|
62 |
+
[tool.poetry.group.pybullet]
|
63 |
+
optional = true
|
64 |
+
[tool.poetry.group.pybullet.dependencies]
|
65 |
+
pybullet = "3.1.8"
|
66 |
+
|
67 |
+
[tool.poetry.group.procgen]
|
68 |
+
optional = true
|
69 |
+
[tool.poetry.group.procgen.dependencies]
|
70 |
+
procgen = "^0.10.7"
|
71 |
+
|
72 |
+
[tool.poetry.group.pytest]
|
73 |
+
optional = true
|
74 |
+
[tool.poetry.group.pytest.dependencies]
|
75 |
+
pytest = "^7.1.3"
|
76 |
+
|
77 |
+
[tool.poetry.group.mujoco]
|
78 |
+
optional = true
|
79 |
+
[tool.poetry.group.mujoco.dependencies]
|
80 |
+
mujoco = "^2.2"
|
81 |
+
imageio = "^2.14.1"
|
82 |
+
|
83 |
+
[tool.poetry.group.mujoco_py]
|
84 |
+
optional = true
|
85 |
+
[tool.poetry.group.mujoco_py.dependencies]
|
86 |
+
free-mujoco-py = "^2.1.6"
|
87 |
+
|
88 |
+
[tool.poetry.group.docs]
|
89 |
+
optional = true
|
90 |
+
[tool.poetry.group.docs.dependencies]
|
91 |
+
mkdocs-material = "^8.4.3"
|
92 |
+
markdown-include = "^0.7.0"
|
93 |
+
|
94 |
+
[tool.poetry.group.jax]
|
95 |
+
optional = true
|
96 |
+
[tool.poetry.group.jax.dependencies]
|
97 |
+
jax = "^0.3.17"
|
98 |
+
jaxlib = "^0.3.15"
|
99 |
+
flax = "^0.6.0"
|
100 |
+
|
101 |
+
[tool.poetry.group.optuna]
|
102 |
+
optional = true
|
103 |
+
[tool.poetry.group.optuna.dependencies]
|
104 |
+
optuna = "^3.0.1"
|
105 |
+
optuna-dashboard = "^0.7.2"
|
106 |
+
rich = "<12.0"
|
107 |
+
|
108 |
+
[tool.poetry.group.envpool]
|
109 |
+
optional = true
|
110 |
+
[tool.poetry.group.envpool.dependencies]
|
111 |
+
envpool = "^0.8.1"
|
112 |
+
|
113 |
+
[tool.poetry.group.pettingzoo]
|
114 |
+
optional = true
|
115 |
+
[tool.poetry.group.pettingzoo.dependencies]
|
116 |
+
PettingZoo = "1.18.1"
|
117 |
+
SuperSuit = "3.4.0"
|
118 |
+
multi-agent-ale-py = "0.1.11"
|
119 |
+
|
120 |
+
[tool.poetry.group.cloud]
|
121 |
+
optional = true
|
122 |
+
[tool.poetry.group.cloud.dependencies]
|
123 |
+
boto3 = "^1.24.70"
|
124 |
+
awscli = "^1.25.71"
|
125 |
+
|
126 |
+
[tool.poetry.group.isaacgym]
|
127 |
+
optional = true
|
128 |
+
[tool.poetry.group.isaacgym.dependencies]
|
129 |
+
isaacgymenvs = {git = "https://github.com/vwxyzjn/IsaacGymEnvs.git", rev = "poetry"}
|
130 |
+
isaacgym = {path = "cleanrl/ppo_continuous_action_isaacgym/isaacgym", develop = true}
|
131 |
+
|
132 |
+
[tool.poetry.group.dm_control]
|
133 |
+
optional = true
|
134 |
+
[tool.poetry.group.dm_control.dependencies]
|
135 |
+
shimmy = "^0.1.0"
|
136 |
+
dm-control = "^1.0.8"
|
137 |
+
mujoco = "^2.2"
|
138 |
+
|
139 |
+
[build-system]
|
140 |
+
requires = ["poetry-core"]
|
141 |
+
build-backend = "poetry.core.masonry.api"
|
142 |
+
|
143 |
+
[tool.poetry.extras]
|
144 |
+
atari = ["ale-py", "AutoROM", "opencv-python"]
|
145 |
+
pybullet = ["pybullet"]
|
146 |
+
procgen = ["procgen"]
|
147 |
+
plot = ["pandas", "seaborn"]
|
148 |
+
pytest = ["pytest"]
|
149 |
+
mujoco = ["mujoco", "imageio"]
|
150 |
+
mujoco_py = ["free-mujoco-py"]
|
151 |
+
jax = ["jax", "jaxlib", "flax"]
|
152 |
+
docs = ["mkdocs-material", "markdown-include"]
|
153 |
+
envpool = ["envpool"]
|
154 |
+
optuna = ["optuna", "optuna-dashboard", "rich"]
|
155 |
+
pettingzoo = ["PettingZoo", "SuperSuit", "multi-agent-ale-py"]
|
156 |
+
cloud = ["boto3", "awscli"]
|
157 |
+
dm_control = ["shimmy", "dm-control", "mujoco"]
|
158 |
+
|
159 |
+
# dependencies for algorithm variant (useful when you want to run a specific algorithm)
|
160 |
+
dqn = []
|
161 |
+
dqn_atari = ["ale-py", "AutoROM", "opencv-python"]
|
162 |
+
dqn_jax = ["jax", "jaxlib", "flax"]
|
163 |
+
dqn_atari_jax = [
|
164 |
+
"ale-py", "AutoROM", "opencv-python", # atari
|
165 |
+
"jax", "jaxlib", "flax" # jax
|
166 |
+
]
|
167 |
+
c51 = []
|
168 |
+
c51_atari = ["ale-py", "AutoROM", "opencv-python"]
|
169 |
+
c51_jax = ["jax", "jaxlib", "flax"]
|
170 |
+
c51_atari_jax = [
|
171 |
+
"ale-py", "AutoROM", "opencv-python", # atari
|
172 |
+
"jax", "jaxlib", "flax" # jax
|
173 |
+
]
|
174 |
+
ppo_atari_envpool_xla_jax_scan = [
|
175 |
+
"ale-py", "AutoROM", "opencv-python", # atari
|
176 |
+
"jax", "jaxlib", "flax", # jax
|
177 |
+
"envpool", # envpool
|
178 |
+
]
|
replay.mp4
ADDED
Binary file (28.9 kB). View file
|
|
sebulba_ppo_envpool.cleanrl_model
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:f2943bd34ada8a809a71dbfeee42aad0015e6d5c71b8d59a002a6b089199f5f7
|
3 |
+
size 4378341
|
sebulba_ppo_envpool.py
ADDED
@@ -0,0 +1,889 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""
|
2 |
+
0. multi-threaded actor
|
3 |
+
python sebulba_ppo_envpool.py --actor-device-ids 0 --num-actor-threads 2 --learner-device-ids 1 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
4 |
+
python sebulba_ppo_envpool.py --actor-device-ids 0 --learner-device-ids 1 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
5 |
+
|
6 |
+
🔥 core settings:
|
7 |
+
|
8 |
+
* test throughput
|
9 |
+
* python sebulba_ppo_envpool.py --exp-name sebula_thpt_a0_l1_timeout --actor-device-ids 0 --learner-device-ids 1 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
10 |
+
* python sebulba_ppo_envpool.py --exp-name sebula_thpt_a0_l12_timeout --actor-device-ids 0 --learner-device-ids 1 2 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
11 |
+
* this will help us diagnose the throughput issue
|
12 |
+
* python sebulba_ppo_envpool.py --exp-name sebula_thpt_a0_l1 --actor-device-ids 0 --learner-device-ids 1 --profile --total-timesteps 500000 --track
|
13 |
+
* python sebulba_ppo_envpool.py --exp-name sebula_thpt_a0_l12 --actor-device-ids 0 --learner-device-ids 1 2 --profile --total-timesteps 500000 --track
|
14 |
+
* python sebulba_ppo_envpool.py --actor-device-ids 0 --learner-device-ids 1 2 3 4 --num-actor-threads 2 --track
|
15 |
+
* Best performance so far
|
16 |
+
* python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l01_rollout_is_faster --actor-device-ids 0 --learner-device-ids 0 1 --total-timesteps 500000 --track
|
17 |
+
* python sebulba_ppo_envpool.py --actor-device-ids 0 --learner-device-ids 1 2 3 4 --params-queue-timeout 0.02 --track
|
18 |
+
|
19 |
+
# 1. rollout is faster than training
|
20 |
+
|
21 |
+
## throughput
|
22 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_thpt_rollout_is_faster --actor-device-ids 0 --learner-device-ids 1 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
23 |
+
|
24 |
+
## shared: actor on GPU0 and learner on GPU0
|
25 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_1gpu_rollout_is_faster --actor-device-ids 0 --learner-device-ids 0 --total-timesteps 500000 --track
|
26 |
+
|
27 |
+
## separate: actor on GPU0 and learner on GPU1
|
28 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l1_rollout_is_faster --actor-device-ids 0 --learner-device-ids 1 --total-timesteps 500000 --track
|
29 |
+
|
30 |
+
## shared: actor on GPU0 and learner on GPU0,1
|
31 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l01_rollout_is_faster --actor-device-ids 0 --learner-device-ids 0 1 --total-timesteps 500000 --track
|
32 |
+
|
33 |
+
## separate: actor on GPU0 and learner on GPU1,2
|
34 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l12_rollout_is_faster --actor-device-ids 0 --learner-device-ids 1 2 --total-timesteps 500000 --track
|
35 |
+
|
36 |
+
|
37 |
+
# 1.1 rollout is faster than training w/ timeout
|
38 |
+
|
39 |
+
## shared: actor on GPU0 and learner on GPU0
|
40 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_1gpu_rollout_is_faster_timeout --actor-device-ids 0 --learner-device-ids 0 --params-queue-timeout 0.02 --total-timesteps 500000 --track
|
41 |
+
|
42 |
+
## separate: actor on GPU0 and learner on GPU1
|
43 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l1_rollout_is_faster_timeout --actor-device-ids 0 --learner-device-ids 1 --params-queue-timeout 0.02 --total-timesteps 500000 --track
|
44 |
+
|
45 |
+
## shared: actor on GPU0 and learner on GPU0,1
|
46 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l01_rollout_is_faster_timeout --actor-device-ids 0 --learner-device-ids 0 1 --params-queue-timeout 0.02 --total-timesteps 500000 --track
|
47 |
+
|
48 |
+
## separate: actor on GPU0 and learner on GPU1,2
|
49 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l12_rollout_is_faster_timeout --actor-device-ids 0 --learner-device-ids 1 2 --params-queue-timeout 0.02 --total-timesteps 500000 --track
|
50 |
+
|
51 |
+
# 1.2. rollout is much faster than training w/ timeout
|
52 |
+
|
53 |
+
## throughput
|
54 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_thpt_rollout_is_much_faster_timeout --actor-device-ids 0 --learner-device-ids 1 --update-epochs 8 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
55 |
+
|
56 |
+
## shared: actor on GPU0 and learner on GPU0,1
|
57 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l01_rollout_is_much_faster_timeout --actor-device-ids 0 --learner-device-ids 0 1 --update-epochs 8 --params-queue-timeout 0.02 --total-timesteps 500000 --track
|
58 |
+
|
59 |
+
## separate: actor on GPU0 and learner on GPU1,2
|
60 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l12_rollout_is_much_faster_timeout --actor-device-ids 0 --learner-device-ids 1 2 --update-epochs 8 --params-queue-timeout 0.02 --total-timesteps 500000 --track
|
61 |
+
|
62 |
+
# 2. training is faster than rollout
|
63 |
+
|
64 |
+
## throughput
|
65 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_thpt_training_is_faster --update-epochs 1 --async-batch-size 64 --actor-device-ids 0 --learner-device-ids 1 --params-queue-timeout 0.02 --profile --test-actor-learner-throughput --total-timesteps 500000 --track
|
66 |
+
|
67 |
+
## shared: actor on GPU0 and learner on GPU0
|
68 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_1gpu_training_is_faster --update-epochs 1 --async-batch-size 64 --actor-device-ids 0 --learner-device-ids 0 --total-timesteps 500000 --track
|
69 |
+
|
70 |
+
## separate: actor on GPU0 and learner on GPU1
|
71 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l1_training_is_faster --update-epochs 1 --async-batch-size 64 --actor-device-ids 0 --learner-device-ids 1 --total-timesteps 500000 --track
|
72 |
+
|
73 |
+
## shared: actor on GPU0 and learner on GPU0,1
|
74 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l01_training_is_faster --update-epochs 1 --async-batch-size 64 --actor-device-ids 0 --learner-device-ids 0 1 --total-timesteps 500000 --track
|
75 |
+
|
76 |
+
## separate: actor on GPU0 and learner on GPU1,2
|
77 |
+
python sebulba_ppo_envpool.py --exp-name sebulba_ppo_envpool_a0_l12_training_is_faster --update-epochs 1 --async-batch-size 64 --actor-device-ids 0 --learner-device-ids 1 2 --total-timesteps 500000 --track
|
78 |
+
|
79 |
+
"""
|
80 |
+
# docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_atari_envpool_async_jax_scan_impalanet_machadopy
|
81 |
+
# https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/
|
82 |
+
import argparse
|
83 |
+
import os
|
84 |
+
import random
|
85 |
+
import time
|
86 |
+
import uuid
|
87 |
+
from collections import deque
|
88 |
+
from distutils.util import strtobool
|
89 |
+
from functools import partial
|
90 |
+
from typing import Sequence
|
91 |
+
|
92 |
+
os.environ[
|
93 |
+
"XLA_PYTHON_CLIENT_MEM_FRACTION"
|
94 |
+
] = "0.6" # see https://github.com/google/jax/discussions/6332#discussioncomment-1279991
|
95 |
+
os.environ["XLA_FLAGS"] = "--xla_cpu_multi_thread_eigen=false " "intra_op_parallelism_threads=1"
|
96 |
+
import multiprocessing as mp
|
97 |
+
import queue
|
98 |
+
import threading
|
99 |
+
|
100 |
+
import envpool
|
101 |
+
import flax
|
102 |
+
import flax.linen as nn
|
103 |
+
import gym
|
104 |
+
import jax
|
105 |
+
import jax.numpy as jnp
|
106 |
+
import numpy as np
|
107 |
+
import optax
|
108 |
+
from flax.linen.initializers import constant, orthogonal
|
109 |
+
from flax.training.train_state import TrainState
|
110 |
+
from torch.utils.tensorboard import SummaryWriter
|
111 |
+
|
112 |
+
|
113 |
+
def parse_args():
|
114 |
+
# fmt: off
|
115 |
+
parser = argparse.ArgumentParser()
|
116 |
+
parser.add_argument("--exp-name", type=str, default=os.path.basename(__file__).rstrip(".py"),
|
117 |
+
help="the name of this experiment")
|
118 |
+
parser.add_argument("--seed", type=int, default=1,
|
119 |
+
help="seed of the experiment")
|
120 |
+
parser.add_argument("--torch-deterministic", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
|
121 |
+
help="if toggled, `torch.backends.cudnn.deterministic=False`")
|
122 |
+
parser.add_argument("--cuda", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
|
123 |
+
help="if toggled, cuda will be enabled by default")
|
124 |
+
parser.add_argument("--track", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
|
125 |
+
help="if toggled, this experiment will be tracked with Weights and Biases")
|
126 |
+
parser.add_argument("--wandb-project-name", type=str, default="cleanRL",
|
127 |
+
help="the wandb's project name")
|
128 |
+
parser.add_argument("--wandb-entity", type=str, default=None,
|
129 |
+
help="the entity (team) of wandb's project")
|
130 |
+
parser.add_argument("--capture-video", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
|
131 |
+
help="weather to capture videos of the agent performances (check out `videos` folder)")
|
132 |
+
parser.add_argument("--save-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
|
133 |
+
help="whether to save model into the `runs/{run_name}` folder")
|
134 |
+
parser.add_argument("--upload-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
|
135 |
+
help="whether to upload the saved model to huggingface")
|
136 |
+
parser.add_argument("--hf-entity", type=str, default="",
|
137 |
+
help="the user or org name of the model repository from the Hugging Face Hub")
|
138 |
+
|
139 |
+
# Algorithm specific arguments
|
140 |
+
parser.add_argument("--env-id", type=str, default="Breakout-v5",
|
141 |
+
help="the id of the environment")
|
142 |
+
parser.add_argument("--total-timesteps", type=int, default=50000000,
|
143 |
+
help="total timesteps of the experiments")
|
144 |
+
parser.add_argument("--learning-rate", type=float, default=2.5e-4,
|
145 |
+
help="the learning rate of the optimizer")
|
146 |
+
parser.add_argument("--num-envs", type=int, default=64,
|
147 |
+
help="the number of parallel game environments")
|
148 |
+
parser.add_argument("--async-batch-size", type=int, default=16,
|
149 |
+
help="the envpool's batch size in the async mode")
|
150 |
+
parser.add_argument("--num-steps", type=int, default=128,
|
151 |
+
help="the number of steps to run in each environment per policy rollout")
|
152 |
+
parser.add_argument("--anneal-lr", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
|
153 |
+
help="Toggle learning rate annealing for policy and value networks")
|
154 |
+
parser.add_argument("--gamma", type=float, default=0.99,
|
155 |
+
help="the discount factor gamma")
|
156 |
+
parser.add_argument("--gae-lambda", type=float, default=0.95,
|
157 |
+
help="the lambda for the general advantage estimation")
|
158 |
+
parser.add_argument("--num-minibatches", type=int, default=4,
|
159 |
+
help="the number of mini-batches")
|
160 |
+
parser.add_argument("--update-epochs", type=int, default=4,
|
161 |
+
help="the K epochs to update the policy")
|
162 |
+
parser.add_argument("--norm-adv", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
|
163 |
+
help="Toggles advantages normalization")
|
164 |
+
parser.add_argument("--clip-coef", type=float, default=0.1,
|
165 |
+
help="the surrogate clipping coefficient")
|
166 |
+
parser.add_argument("--ent-coef", type=float, default=0.01,
|
167 |
+
help="coefficient of the entropy")
|
168 |
+
parser.add_argument("--vf-coef", type=float, default=0.5,
|
169 |
+
help="coefficient of the value function")
|
170 |
+
parser.add_argument("--max-grad-norm", type=float, default=0.5,
|
171 |
+
help="the maximum norm for the gradient clipping")
|
172 |
+
parser.add_argument("--target-kl", type=float, default=None,
|
173 |
+
help="the target KL divergence threshold")
|
174 |
+
|
175 |
+
parser.add_argument("--actor-device-ids", type=int, nargs="+", default=[0], # type is actually List[int]
|
176 |
+
help="the device ids that actor workers will use")
|
177 |
+
parser.add_argument("--learner-device-ids", type=int, nargs="+", default=[0], # type is actually List[int]
|
178 |
+
help="the device ids that actor workers will use")
|
179 |
+
parser.add_argument("--num-actor-threads", type=int, default=1,
|
180 |
+
help="the number of actor threads")
|
181 |
+
parser.add_argument("--profile", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
|
182 |
+
help="whether to call block_until_ready() for profiling")
|
183 |
+
parser.add_argument("--test-actor-learner-throughput", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
|
184 |
+
help="whether to test actor-learner throughput by removing the actor-learner communication")
|
185 |
+
parser.add_argument("--params-queue-timeout", type=float, default=None,
|
186 |
+
help="the timeout for the `params_queue.get()` operation in the actor thread to pull params;" + \
|
187 |
+
"by default it's `None`; if you set a timeout, it will likely make the actor run faster but will introduce some side effects," + \
|
188 |
+
"such as the actor will not be able to pull the latest params from the learner and will use the old params instead")
|
189 |
+
args = parser.parse_args()
|
190 |
+
args.batch_size = int(args.num_envs * args.num_steps)
|
191 |
+
args.minibatch_size = int(args.batch_size // args.num_minibatches)
|
192 |
+
args.num_updates = args.total_timesteps // args.batch_size
|
193 |
+
args.async_update = int(args.num_envs / args.async_batch_size)
|
194 |
+
assert len(args.actor_device_ids) == 1, "only 1 actor_device_ids is supported now"
|
195 |
+
# fmt: on
|
196 |
+
return args
|
197 |
+
|
198 |
+
|
199 |
+
LEARNER_WARMUP_TIME = 10 # seconds
|
200 |
+
|
201 |
+
|
202 |
+
def make_env(env_id, seed, num_envs, async_batch_size=1, num_threads=None, thread_affinity_offset=-1):
|
203 |
+
def thunk():
|
204 |
+
envs = envpool.make(
|
205 |
+
env_id,
|
206 |
+
env_type="gym",
|
207 |
+
num_envs=num_envs,
|
208 |
+
num_threads=num_threads if num_threads is not None else async_batch_size,
|
209 |
+
thread_affinity_offset=thread_affinity_offset,
|
210 |
+
batch_size=async_batch_size,
|
211 |
+
episodic_life=False, # Machado et al. 2017 (Revisitng ALE: Eval protocols) p. 6
|
212 |
+
repeat_action_probability=0.25, # Machado et al. 2017 (Revisitng ALE: Eval protocols) p. 12
|
213 |
+
noop_max=1, # Machado et al. 2017 (Revisitng ALE: Eval protocols) p. 12 (no-op is deprecated in favor of sticky action, right?)
|
214 |
+
full_action_space=True, # Machado et al. 2017 (Revisitng ALE: Eval protocols) Tab. 5
|
215 |
+
max_episode_steps=int(108000 / 4), # Hessel et al. 2018 (Rainbow DQN), Table 3, Max frames per episode
|
216 |
+
reward_clip=True,
|
217 |
+
seed=seed,
|
218 |
+
)
|
219 |
+
envs.num_envs = num_envs
|
220 |
+
envs.single_action_space = envs.action_space
|
221 |
+
envs.single_observation_space = envs.observation_space
|
222 |
+
envs.is_vector_env = True
|
223 |
+
return envs
|
224 |
+
|
225 |
+
return thunk
|
226 |
+
|
227 |
+
|
228 |
+
class ResidualBlock(nn.Module):
|
229 |
+
channels: int
|
230 |
+
|
231 |
+
@nn.compact
|
232 |
+
def __call__(self, x):
|
233 |
+
inputs = x
|
234 |
+
x = nn.relu(x)
|
235 |
+
x = nn.Conv(
|
236 |
+
self.channels,
|
237 |
+
kernel_size=(3, 3),
|
238 |
+
)(x)
|
239 |
+
x = nn.relu(x)
|
240 |
+
x = nn.Conv(
|
241 |
+
self.channels,
|
242 |
+
kernel_size=(3, 3),
|
243 |
+
)(x)
|
244 |
+
return x + inputs
|
245 |
+
|
246 |
+
|
247 |
+
class ConvSequence(nn.Module):
|
248 |
+
channels: int
|
249 |
+
|
250 |
+
@nn.compact
|
251 |
+
def __call__(self, x):
|
252 |
+
x = nn.Conv(
|
253 |
+
self.channels,
|
254 |
+
kernel_size=(3, 3),
|
255 |
+
)(x)
|
256 |
+
x = nn.max_pool(x, window_shape=(3, 3), strides=(2, 2), padding="SAME")
|
257 |
+
x = ResidualBlock(self.channels)(x)
|
258 |
+
x = ResidualBlock(self.channels)(x)
|
259 |
+
return x
|
260 |
+
|
261 |
+
|
262 |
+
class Network(nn.Module):
|
263 |
+
channelss: Sequence[int] = (16, 32, 32)
|
264 |
+
|
265 |
+
@nn.compact
|
266 |
+
def __call__(self, x):
|
267 |
+
x = jnp.transpose(x, (0, 2, 3, 1))
|
268 |
+
x = x / (255.0)
|
269 |
+
for channels in self.channelss:
|
270 |
+
x = ConvSequence(channels)(x)
|
271 |
+
x = nn.relu(x)
|
272 |
+
x = x.reshape((x.shape[0], -1))
|
273 |
+
x = nn.Dense(256, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(x)
|
274 |
+
x = nn.relu(x)
|
275 |
+
return x
|
276 |
+
|
277 |
+
|
278 |
+
class Critic(nn.Module):
|
279 |
+
@nn.compact
|
280 |
+
def __call__(self, x):
|
281 |
+
return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
|
282 |
+
|
283 |
+
|
284 |
+
class Actor(nn.Module):
|
285 |
+
action_dim: int
|
286 |
+
|
287 |
+
@nn.compact
|
288 |
+
def __call__(self, x):
|
289 |
+
return nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
|
290 |
+
|
291 |
+
|
292 |
+
@flax.struct.dataclass
|
293 |
+
class AgentParams:
|
294 |
+
network_params: flax.core.FrozenDict
|
295 |
+
actor_params: flax.core.FrozenDict
|
296 |
+
critic_params: flax.core.FrozenDict
|
297 |
+
|
298 |
+
|
299 |
+
@partial(jax.jit, static_argnums=(3))
|
300 |
+
def get_action_and_value(
|
301 |
+
params: TrainState,
|
302 |
+
next_obs: np.ndarray,
|
303 |
+
key: jax.random.PRNGKey,
|
304 |
+
action_dim: int,
|
305 |
+
):
|
306 |
+
hidden = Network().apply(params.network_params, next_obs)
|
307 |
+
logits = Actor(action_dim).apply(params.actor_params, hidden)
|
308 |
+
# sample action: Gumbel-softmax trick
|
309 |
+
# see https://stats.stackexchange.com/questions/359442/sampling-from-a-categorical-distribution
|
310 |
+
key, subkey = jax.random.split(key)
|
311 |
+
u = jax.random.uniform(subkey, shape=logits.shape)
|
312 |
+
action = jnp.argmax(logits - jnp.log(-jnp.log(u)), axis=1)
|
313 |
+
logprob = jax.nn.log_softmax(logits)[jnp.arange(action.shape[0]), action]
|
314 |
+
value = Critic().apply(params.critic_params, hidden)
|
315 |
+
return action, logprob, value.squeeze(), key
|
316 |
+
|
317 |
+
|
318 |
+
@jax.jit
|
319 |
+
def prepare_data(
|
320 |
+
obs: list,
|
321 |
+
dones: list,
|
322 |
+
values: list,
|
323 |
+
actions: list,
|
324 |
+
logprobs: list,
|
325 |
+
env_ids: list,
|
326 |
+
rewards: list,
|
327 |
+
):
|
328 |
+
obs = jnp.asarray(obs)
|
329 |
+
dones = jnp.asarray(dones)
|
330 |
+
values = jnp.asarray(values)
|
331 |
+
actions = jnp.asarray(actions)
|
332 |
+
logprobs = jnp.asarray(logprobs)
|
333 |
+
env_ids = jnp.asarray(env_ids)
|
334 |
+
rewards = jnp.asarray(rewards)
|
335 |
+
|
336 |
+
# TODO: in an unlikely event, one of the envs might have not stepped at all, which may results in unexpected behavior
|
337 |
+
T, B = env_ids.shape
|
338 |
+
index_ranges = jnp.arange(T * B, dtype=jnp.int32)
|
339 |
+
next_index_ranges = jnp.zeros_like(index_ranges, dtype=jnp.int32)
|
340 |
+
last_env_ids = jnp.zeros(args.num_envs, dtype=jnp.int32) - 1
|
341 |
+
|
342 |
+
def f(carry, x):
|
343 |
+
last_env_ids, next_index_ranges = carry
|
344 |
+
env_id, index_range = x
|
345 |
+
next_index_ranges = next_index_ranges.at[last_env_ids[env_id]].set(
|
346 |
+
jnp.where(last_env_ids[env_id] != -1, index_range, next_index_ranges[last_env_ids[env_id]])
|
347 |
+
)
|
348 |
+
last_env_ids = last_env_ids.at[env_id].set(index_range)
|
349 |
+
return (last_env_ids, next_index_ranges), None
|
350 |
+
|
351 |
+
(last_env_ids, next_index_ranges), _ = jax.lax.scan(
|
352 |
+
f,
|
353 |
+
(last_env_ids, next_index_ranges),
|
354 |
+
(env_ids.reshape(-1), index_ranges),
|
355 |
+
)
|
356 |
+
|
357 |
+
# rewards is off by one time step
|
358 |
+
rewards = rewards.reshape(-1)[next_index_ranges].reshape((args.num_steps) * args.async_update, args.async_batch_size)
|
359 |
+
advantages, returns, _, final_env_ids = compute_gae(env_ids, rewards, values, dones)
|
360 |
+
# b_inds = jnp.nonzero(final_env_ids.reshape(-1), size=(args.num_steps) * args.async_update * args.async_batch_size)[0] # useful for debugging
|
361 |
+
b_obs = obs.reshape((-1,) + obs.shape[2:])
|
362 |
+
b_actions = actions.reshape(-1)
|
363 |
+
b_logprobs = logprobs.reshape(-1)
|
364 |
+
b_advantages = advantages.reshape(-1)
|
365 |
+
b_returns = returns.reshape(-1)
|
366 |
+
return b_obs, b_actions, b_logprobs, b_advantages, b_returns
|
367 |
+
|
368 |
+
|
369 |
+
def rollout(
|
370 |
+
i,
|
371 |
+
num_threads, # =None,
|
372 |
+
thread_affinity_offset, # =-1,
|
373 |
+
key: jax.random.PRNGKey,
|
374 |
+
args,
|
375 |
+
rollout_queue,
|
376 |
+
params_queue: queue.Queue,
|
377 |
+
writer,
|
378 |
+
learner_devices,
|
379 |
+
):
|
380 |
+
envs = make_env(args.env_id, args.seed, args.num_envs, args.async_batch_size, num_threads, thread_affinity_offset)()
|
381 |
+
len_actor_device_ids = len(args.actor_device_ids)
|
382 |
+
global_step = 0
|
383 |
+
# TRY NOT TO MODIFY: start the game
|
384 |
+
start_time = time.time()
|
385 |
+
|
386 |
+
# put data in the last index
|
387 |
+
episode_returns = np.zeros((args.num_envs,), dtype=np.float32)
|
388 |
+
returned_episode_returns = np.zeros((args.num_envs,), dtype=np.float32)
|
389 |
+
episode_lengths = np.zeros((args.num_envs,), dtype=np.float32)
|
390 |
+
returned_episode_lengths = np.zeros((args.num_envs,), dtype=np.float32)
|
391 |
+
envs.async_reset()
|
392 |
+
|
393 |
+
params_queue_get_time = deque(maxlen=10)
|
394 |
+
rollout_time = deque(maxlen=10)
|
395 |
+
data_transfer_time = deque(maxlen=10)
|
396 |
+
rollout_queue_put_time = deque(maxlen=10)
|
397 |
+
params_timeout_count = 0
|
398 |
+
for update in range(1, args.num_updates + 2):
|
399 |
+
update_time_start = time.time()
|
400 |
+
obs = []
|
401 |
+
dones = []
|
402 |
+
actions = []
|
403 |
+
logprobs = []
|
404 |
+
values = []
|
405 |
+
env_ids = []
|
406 |
+
rewards = []
|
407 |
+
truncations = []
|
408 |
+
terminations = []
|
409 |
+
env_recv_time = 0
|
410 |
+
inference_time = 0
|
411 |
+
storage_time = 0
|
412 |
+
env_send_time = 0
|
413 |
+
|
414 |
+
# NOTE: This is a major difference from the sync version:
|
415 |
+
# at the end of the rollout phase, the sync version will have the next observation
|
416 |
+
# ready for the value bootstrap, but the async version will not have it.
|
417 |
+
# for this reason we do `num_steps + 1`` to get the extra states for value bootstrapping.
|
418 |
+
# but note that the extra states are not used for the loss computation in the next iteration,
|
419 |
+
# while the sync version will use the extra state for the loss computation.
|
420 |
+
params_queue_get_time_start = time.time()
|
421 |
+
try:
|
422 |
+
params = params_queue.get(timeout=args.params_queue_timeout)
|
423 |
+
except queue.Empty:
|
424 |
+
# print("params_queue.get timeout triggered")
|
425 |
+
params_timeout_count += 1
|
426 |
+
params_queue_get_time.append(time.time() - params_queue_get_time_start)
|
427 |
+
writer.add_scalar("stats/params_queue_get_time", np.mean(params_queue_get_time), global_step)
|
428 |
+
writer.add_scalar("stats/params_queue_timeout_count", params_timeout_count, global_step)
|
429 |
+
rollout_time_start = time.time()
|
430 |
+
for _ in range(
|
431 |
+
args.async_update, (args.num_steps + 1) * args.async_update
|
432 |
+
): # num_steps + 1 to get the states for value bootstrapping.
|
433 |
+
env_recv_time_start = time.time()
|
434 |
+
next_obs, next_reward, next_done, info = envs.recv()
|
435 |
+
env_recv_time += time.time() - env_recv_time_start
|
436 |
+
global_step += len(next_done) * args.num_actor_threads * len_actor_device_ids
|
437 |
+
env_id = info["env_id"]
|
438 |
+
|
439 |
+
inference_time_start = time.time()
|
440 |
+
action, logprob, value, key = get_action_and_value(params, next_obs, key, envs.single_action_space.n)
|
441 |
+
inference_time += time.time() - inference_time_start
|
442 |
+
|
443 |
+
env_send_time_start = time.time()
|
444 |
+
envs.send(np.array(action), env_id)
|
445 |
+
env_send_time += time.time() - env_send_time_start
|
446 |
+
storage_time_start = time.time()
|
447 |
+
obs.append(next_obs)
|
448 |
+
dones.append(next_done)
|
449 |
+
values.append(value)
|
450 |
+
actions.append(action)
|
451 |
+
logprobs.append(logprob)
|
452 |
+
env_ids.append(env_id)
|
453 |
+
rewards.append(next_reward)
|
454 |
+
truncations.append(info["TimeLimit.truncated"])
|
455 |
+
terminations.append(info["terminated"])
|
456 |
+
episode_returns[env_id] += info["reward"]
|
457 |
+
returned_episode_returns[env_id] = np.where(
|
458 |
+
info["terminated"] + info["TimeLimit.truncated"], episode_returns[env_id], returned_episode_returns[env_id]
|
459 |
+
)
|
460 |
+
episode_returns[env_id] *= (1 - info["terminated"]) * (1 - info["TimeLimit.truncated"])
|
461 |
+
episode_lengths[env_id] += 1
|
462 |
+
returned_episode_lengths[env_id] = np.where(
|
463 |
+
info["terminated"] + info["TimeLimit.truncated"], episode_lengths[env_id], returned_episode_lengths[env_id]
|
464 |
+
)
|
465 |
+
episode_lengths[env_id] *= (1 - info["terminated"]) * (1 - info["TimeLimit.truncated"])
|
466 |
+
storage_time += time.time() - storage_time_start
|
467 |
+
if args.profile:
|
468 |
+
action.block_until_ready()
|
469 |
+
rollout_time.append(time.time() - rollout_time_start)
|
470 |
+
writer.add_scalar("stats/rollout_time", np.mean(rollout_time), global_step)
|
471 |
+
|
472 |
+
avg_episodic_return = np.mean(returned_episode_returns)
|
473 |
+
writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
|
474 |
+
writer.add_scalar("charts/avg_episodic_length", np.mean(returned_episode_lengths), global_step)
|
475 |
+
if i == 0:
|
476 |
+
print(f"global_step={global_step}, avg_episodic_return={avg_episodic_return}")
|
477 |
+
print("SPS:", int(global_step / (time.time() - start_time)))
|
478 |
+
writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
|
479 |
+
|
480 |
+
writer.add_scalar("stats/truncations", np.sum(truncations), global_step)
|
481 |
+
writer.add_scalar("stats/terminations", np.sum(terminations), global_step)
|
482 |
+
writer.add_scalar("stats/env_recv_time", env_recv_time, global_step)
|
483 |
+
writer.add_scalar("stats/inference_time", inference_time, global_step)
|
484 |
+
writer.add_scalar("stats/storage_time", storage_time, global_step)
|
485 |
+
writer.add_scalar("stats/env_send_time", env_send_time, global_step)
|
486 |
+
|
487 |
+
data_transfer_time_start = time.time()
|
488 |
+
b_obs, b_actions, b_logprobs, b_advantages, b_returns = prepare_data(
|
489 |
+
obs,
|
490 |
+
dones,
|
491 |
+
values,
|
492 |
+
actions,
|
493 |
+
logprobs,
|
494 |
+
env_ids,
|
495 |
+
rewards,
|
496 |
+
)
|
497 |
+
payload = (
|
498 |
+
global_step,
|
499 |
+
update,
|
500 |
+
jnp.array_split(b_obs, len(learner_devices)),
|
501 |
+
jnp.array_split(b_actions, len(learner_devices)),
|
502 |
+
jnp.array_split(b_logprobs, len(learner_devices)),
|
503 |
+
jnp.array_split(b_advantages, len(learner_devices)),
|
504 |
+
jnp.array_split(b_returns, len(learner_devices)),
|
505 |
+
)
|
506 |
+
if args.profile:
|
507 |
+
payload[2][0].block_until_ready()
|
508 |
+
data_transfer_time.append(time.time() - data_transfer_time_start)
|
509 |
+
writer.add_scalar("stats/data_transfer_time", np.mean(data_transfer_time), global_step)
|
510 |
+
if update == 1 or not args.test_actor_learner_throughput:
|
511 |
+
rollout_queue_put_time_start = time.time()
|
512 |
+
rollout_queue.put(payload)
|
513 |
+
rollout_queue_put_time.append(time.time() - rollout_queue_put_time_start)
|
514 |
+
writer.add_scalar("stats/rollout_queue_put_time", np.mean(rollout_queue_put_time), global_step)
|
515 |
+
|
516 |
+
if update == 1 or update == 2 or update == 3:
|
517 |
+
time.sleep(LEARNER_WARMUP_TIME) # makes sure the actor does to fill the rollout_queue at the get go
|
518 |
+
|
519 |
+
writer.add_scalar(
|
520 |
+
"charts/SPS_update",
|
521 |
+
int(
|
522 |
+
args.num_envs
|
523 |
+
* args.num_steps
|
524 |
+
* args.num_actor_threads
|
525 |
+
* len_actor_device_ids
|
526 |
+
/ (time.time() - update_time_start)
|
527 |
+
),
|
528 |
+
global_step,
|
529 |
+
)
|
530 |
+
|
531 |
+
|
532 |
+
@partial(jax.jit, static_argnums=(3))
|
533 |
+
def get_action_and_value2(
|
534 |
+
params: flax.core.FrozenDict,
|
535 |
+
x: np.ndarray,
|
536 |
+
action: np.ndarray,
|
537 |
+
action_dim: int,
|
538 |
+
):
|
539 |
+
hidden = Network().apply(params.network_params, x)
|
540 |
+
logits = Actor(action_dim).apply(params.actor_params, hidden)
|
541 |
+
logprob = jax.nn.log_softmax(logits)[jnp.arange(action.shape[0]), action]
|
542 |
+
logits = logits - jax.scipy.special.logsumexp(logits, axis=-1, keepdims=True)
|
543 |
+
logits = logits.clip(min=jnp.finfo(logits.dtype).min)
|
544 |
+
p_log_p = logits * jax.nn.softmax(logits)
|
545 |
+
entropy = -p_log_p.sum(-1)
|
546 |
+
value = Critic().apply(params.critic_params, hidden).squeeze()
|
547 |
+
return logprob, entropy, value
|
548 |
+
|
549 |
+
|
550 |
+
@jax.jit
|
551 |
+
def compute_gae(
|
552 |
+
env_ids: np.ndarray,
|
553 |
+
rewards: np.ndarray,
|
554 |
+
values: np.ndarray,
|
555 |
+
dones: np.ndarray,
|
556 |
+
):
|
557 |
+
dones = jnp.asarray(dones)
|
558 |
+
values = jnp.asarray(values)
|
559 |
+
env_ids = jnp.asarray(env_ids)
|
560 |
+
rewards = jnp.asarray(rewards)
|
561 |
+
|
562 |
+
_, B = env_ids.shape
|
563 |
+
final_env_id_checked = jnp.zeros(args.num_envs, jnp.int32) - 1
|
564 |
+
final_env_ids = jnp.zeros(B, jnp.int32)
|
565 |
+
advantages = jnp.zeros(B)
|
566 |
+
lastgaelam = jnp.zeros(args.num_envs)
|
567 |
+
lastdones = jnp.zeros(args.num_envs) + 1
|
568 |
+
lastvalues = jnp.zeros(args.num_envs)
|
569 |
+
|
570 |
+
def compute_gae_once(carry, x):
|
571 |
+
lastvalues, lastdones, advantages, lastgaelam, final_env_ids, final_env_id_checked = carry
|
572 |
+
(
|
573 |
+
done,
|
574 |
+
value,
|
575 |
+
eid,
|
576 |
+
reward,
|
577 |
+
) = x
|
578 |
+
nextnonterminal = 1.0 - lastdones[eid]
|
579 |
+
nextvalues = lastvalues[eid]
|
580 |
+
delta = jnp.where(final_env_id_checked[eid] == -1, 0, reward + args.gamma * nextvalues * nextnonterminal - value)
|
581 |
+
advantages = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam[eid]
|
582 |
+
final_env_ids = jnp.where(final_env_id_checked[eid] == 1, 1, 0)
|
583 |
+
final_env_id_checked = final_env_id_checked.at[eid].set(
|
584 |
+
jnp.where(final_env_id_checked[eid] == -1, 1, final_env_id_checked[eid])
|
585 |
+
)
|
586 |
+
|
587 |
+
# the last_ variables keeps track of the actual `num_steps`
|
588 |
+
lastgaelam = lastgaelam.at[eid].set(advantages)
|
589 |
+
lastdones = lastdones.at[eid].set(done)
|
590 |
+
lastvalues = lastvalues.at[eid].set(value)
|
591 |
+
return (lastvalues, lastdones, advantages, lastgaelam, final_env_ids, final_env_id_checked), (
|
592 |
+
advantages,
|
593 |
+
final_env_ids,
|
594 |
+
)
|
595 |
+
|
596 |
+
(_, _, _, _, final_env_ids, final_env_id_checked), (advantages, final_env_ids) = jax.lax.scan(
|
597 |
+
compute_gae_once,
|
598 |
+
(
|
599 |
+
lastvalues,
|
600 |
+
lastdones,
|
601 |
+
advantages,
|
602 |
+
lastgaelam,
|
603 |
+
final_env_ids,
|
604 |
+
final_env_id_checked,
|
605 |
+
),
|
606 |
+
(
|
607 |
+
dones,
|
608 |
+
values,
|
609 |
+
env_ids,
|
610 |
+
rewards,
|
611 |
+
),
|
612 |
+
reverse=True,
|
613 |
+
)
|
614 |
+
return advantages, advantages + values, final_env_id_checked, final_env_ids
|
615 |
+
|
616 |
+
|
617 |
+
def ppo_loss(params, x, a, logp, mb_advantages, mb_returns, action_dim):
|
618 |
+
newlogprob, entropy, newvalue = get_action_and_value2(params, x, a, action_dim)
|
619 |
+
logratio = newlogprob - logp
|
620 |
+
ratio = jnp.exp(logratio)
|
621 |
+
approx_kl = ((ratio - 1) - logratio).mean()
|
622 |
+
|
623 |
+
if args.norm_adv:
|
624 |
+
mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
|
625 |
+
|
626 |
+
# Policy loss
|
627 |
+
pg_loss1 = -mb_advantages * ratio
|
628 |
+
pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
|
629 |
+
pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
|
630 |
+
|
631 |
+
# Value loss
|
632 |
+
v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
|
633 |
+
|
634 |
+
entropy_loss = entropy.mean()
|
635 |
+
loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
|
636 |
+
return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
|
637 |
+
|
638 |
+
|
639 |
+
@partial(jax.jit, static_argnums=(6))
|
640 |
+
def single_device_update(
|
641 |
+
agent_state: TrainState,
|
642 |
+
b_obs,
|
643 |
+
b_actions,
|
644 |
+
b_logprobs,
|
645 |
+
b_advantages,
|
646 |
+
b_returns,
|
647 |
+
action_dim,
|
648 |
+
key: jax.random.PRNGKey,
|
649 |
+
):
|
650 |
+
ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True)
|
651 |
+
|
652 |
+
def update_epoch(carry, _):
|
653 |
+
agent_state, key = carry
|
654 |
+
key, subkey = jax.random.split(key)
|
655 |
+
|
656 |
+
# taken from: https://github.com/google/brax/blob/main/brax/training/agents/ppo/train.py
|
657 |
+
def convert_data(x: jnp.ndarray):
|
658 |
+
x = jax.random.permutation(subkey, x)
|
659 |
+
x = jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
|
660 |
+
return x
|
661 |
+
|
662 |
+
def update_minibatch(agent_state, minibatch):
|
663 |
+
mb_obs, mb_actions, mb_logprobs, mb_advantages, mb_returns = minibatch
|
664 |
+
(loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
|
665 |
+
agent_state.params,
|
666 |
+
mb_obs,
|
667 |
+
mb_actions,
|
668 |
+
mb_logprobs,
|
669 |
+
mb_advantages,
|
670 |
+
mb_returns,
|
671 |
+
action_dim,
|
672 |
+
)
|
673 |
+
grads = jax.lax.pmean(grads, axis_name="devices")
|
674 |
+
agent_state = agent_state.apply_gradients(grads=grads)
|
675 |
+
return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
|
676 |
+
|
677 |
+
agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
|
678 |
+
update_minibatch,
|
679 |
+
agent_state,
|
680 |
+
(
|
681 |
+
convert_data(b_obs),
|
682 |
+
convert_data(b_actions),
|
683 |
+
convert_data(b_logprobs),
|
684 |
+
convert_data(b_advantages),
|
685 |
+
convert_data(b_returns),
|
686 |
+
),
|
687 |
+
)
|
688 |
+
return (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
|
689 |
+
|
690 |
+
(agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, _) = jax.lax.scan(
|
691 |
+
update_epoch, (agent_state, key), (), length=args.update_epochs
|
692 |
+
)
|
693 |
+
return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
|
694 |
+
|
695 |
+
|
696 |
+
if __name__ == "__main__":
|
697 |
+
devices = jax.devices("gpu")
|
698 |
+
args = parse_args()
|
699 |
+
run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{uuid.uuid4()}"
|
700 |
+
if args.track:
|
701 |
+
import wandb
|
702 |
+
|
703 |
+
wandb.init(
|
704 |
+
project=args.wandb_project_name,
|
705 |
+
entity=args.wandb_entity,
|
706 |
+
sync_tensorboard=True,
|
707 |
+
config=vars(args),
|
708 |
+
name=run_name,
|
709 |
+
monitor_gym=True,
|
710 |
+
save_code=True,
|
711 |
+
)
|
712 |
+
writer = SummaryWriter(f"runs/{run_name}")
|
713 |
+
writer.add_text(
|
714 |
+
"hyperparameters",
|
715 |
+
"|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
|
716 |
+
)
|
717 |
+
|
718 |
+
# TRY NOT TO MODIFY: seeding
|
719 |
+
random.seed(args.seed)
|
720 |
+
np.random.seed(args.seed)
|
721 |
+
key = jax.random.PRNGKey(args.seed)
|
722 |
+
key, network_key, actor_key, critic_key = jax.random.split(key, 4)
|
723 |
+
|
724 |
+
# env setup
|
725 |
+
envs = make_env(args.env_id, args.seed, args.num_envs, args.async_batch_size)()
|
726 |
+
assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
|
727 |
+
|
728 |
+
def linear_schedule(count):
|
729 |
+
# anneal learning rate linearly after one training iteration which contains
|
730 |
+
# (args.num_minibatches * args.update_epochs) gradient updates
|
731 |
+
frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_updates
|
732 |
+
return args.learning_rate * frac
|
733 |
+
|
734 |
+
network = Network()
|
735 |
+
actor = Actor(action_dim=envs.single_action_space.n)
|
736 |
+
critic = Critic()
|
737 |
+
network_params = network.init(network_key, np.array([envs.single_observation_space.sample()]))
|
738 |
+
agent_state = TrainState.create(
|
739 |
+
apply_fn=None,
|
740 |
+
params=AgentParams(
|
741 |
+
network_params,
|
742 |
+
actor.init(actor_key, network.apply(network_params, np.array([envs.single_observation_space.sample()]))),
|
743 |
+
critic.init(critic_key, network.apply(network_params, np.array([envs.single_observation_space.sample()]))),
|
744 |
+
),
|
745 |
+
tx=optax.chain(
|
746 |
+
optax.clip_by_global_norm(args.max_grad_norm),
|
747 |
+
optax.inject_hyperparams(optax.adam)(
|
748 |
+
learning_rate=linear_schedule if args.anneal_lr else args.learning_rate, eps=1e-5
|
749 |
+
),
|
750 |
+
),
|
751 |
+
)
|
752 |
+
learner_devices = [devices[d_id] for d_id in args.learner_device_ids]
|
753 |
+
actor_devices = [devices[d_id] for d_id in args.actor_device_ids]
|
754 |
+
agent_state = flax.jax_utils.replicate(agent_state, devices=learner_devices)
|
755 |
+
|
756 |
+
multi_device_update = jax.pmap(
|
757 |
+
single_device_update,
|
758 |
+
axis_name="devices",
|
759 |
+
devices=learner_devices,
|
760 |
+
in_axes=(0, 0, 0, 0, 0, 0, None, None),
|
761 |
+
out_axes=(0, 0, 0, 0, 0, 0, None),
|
762 |
+
static_broadcasted_argnums=(6),
|
763 |
+
)
|
764 |
+
|
765 |
+
rollout_queue = queue.Queue(maxsize=2)
|
766 |
+
params_queues = []
|
767 |
+
num_cpus = mp.cpu_count()
|
768 |
+
fair_num_cpus = num_cpus // len(args.actor_device_ids)
|
769 |
+
|
770 |
+
class DummyWriter:
|
771 |
+
def add_scalar(self, arg0, arg1, arg3):
|
772 |
+
pass
|
773 |
+
|
774 |
+
# lock = threading.Lock()
|
775 |
+
# AgentParamsStore = namedtuple("AgentParamsStore", ["params", "version"])
|
776 |
+
# agent_params_store = AgentParamsStore(agent_state.params, 0)
|
777 |
+
|
778 |
+
dummy_writer = DummyWriter()
|
779 |
+
for d_idx, d_id in enumerate(args.actor_device_ids):
|
780 |
+
for j in range(args.num_actor_threads):
|
781 |
+
params_queue = queue.Queue(maxsize=2)
|
782 |
+
params_queue.put(jax.device_put(flax.jax_utils.unreplicate(agent_state.params), devices[d_id]))
|
783 |
+
threading.Thread(
|
784 |
+
target=rollout,
|
785 |
+
args=(
|
786 |
+
j,
|
787 |
+
fair_num_cpus if args.num_actor_threads > 1 else None,
|
788 |
+
j * args.num_actor_threads if args.num_actor_threads > 1 else -1,
|
789 |
+
jax.device_put(key, devices[d_id]),
|
790 |
+
args,
|
791 |
+
rollout_queue,
|
792 |
+
params_queue,
|
793 |
+
writer if d_idx == 0 and j == 0 else dummy_writer,
|
794 |
+
learner_devices,
|
795 |
+
),
|
796 |
+
).start()
|
797 |
+
params_queues.append(params_queue)
|
798 |
+
|
799 |
+
rollout_queue_get_time = deque(maxlen=10)
|
800 |
+
learner_update = 0
|
801 |
+
while True:
|
802 |
+
learner_update += 1
|
803 |
+
if learner_update == 1 or not args.test_actor_learner_throughput:
|
804 |
+
rollout_queue_get_time_start = time.time()
|
805 |
+
global_step, update, b_obs, b_actions, b_logprobs, b_advantages, b_returns = rollout_queue.get()
|
806 |
+
rollout_queue_get_time.append(time.time() - rollout_queue_get_time_start)
|
807 |
+
writer.add_scalar("stats/rollout_queue_get_time", np.mean(rollout_queue_get_time), global_step)
|
808 |
+
|
809 |
+
training_time_start = time.time()
|
810 |
+
(agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key) = multi_device_update(
|
811 |
+
agent_state,
|
812 |
+
jax.device_put_sharded(b_obs, learner_devices),
|
813 |
+
jax.device_put_sharded(b_actions, learner_devices),
|
814 |
+
jax.device_put_sharded(b_logprobs, learner_devices),
|
815 |
+
jax.device_put_sharded(b_advantages, learner_devices),
|
816 |
+
jax.device_put_sharded(b_returns, learner_devices),
|
817 |
+
envs.single_action_space.n,
|
818 |
+
key,
|
819 |
+
)
|
820 |
+
if learner_update == 1 or not args.test_actor_learner_throughput:
|
821 |
+
for d_idx, d_id in enumerate(args.actor_device_ids):
|
822 |
+
for j in range(args.num_actor_threads):
|
823 |
+
params_queues[d_idx * args.num_actor_threads + j].put(
|
824 |
+
jax.device_put(flax.jax_utils.unreplicate(agent_state.params), devices[d_id])
|
825 |
+
)
|
826 |
+
if args.profile:
|
827 |
+
v_loss[-1, -1, -1].block_until_ready()
|
828 |
+
writer.add_scalar("stats/training_time", time.time() - training_time_start, global_step)
|
829 |
+
writer.add_scalar("stats/rollout_queue_size", rollout_queue.qsize(), global_step)
|
830 |
+
writer.add_scalar("stats/params_queue_size", params_queue.qsize(), global_step)
|
831 |
+
print(global_step, update, rollout_queue.qsize(), f"training time: {time.time() - training_time_start}s")
|
832 |
+
|
833 |
+
# TRY NOT TO MODIFY: record rewards for plotting purposes
|
834 |
+
writer.add_scalar("charts/learning_rate", agent_state.opt_state[1].hyperparams["learning_rate"][0].item(), global_step)
|
835 |
+
writer.add_scalar("losses/value_loss", v_loss[-1, -1, -1].item(), global_step)
|
836 |
+
writer.add_scalar("losses/policy_loss", pg_loss[-1, -1, -1].item(), global_step)
|
837 |
+
writer.add_scalar("losses/entropy", entropy_loss[-1, -1, -1].item(), global_step)
|
838 |
+
writer.add_scalar("losses/approx_kl", approx_kl[-1, -1, -1].item(), global_step)
|
839 |
+
writer.add_scalar("losses/loss", loss[-1, -1, -1].item(), global_step)
|
840 |
+
if update > args.num_updates:
|
841 |
+
break
|
842 |
+
|
843 |
+
if args.save_model:
|
844 |
+
agent_state = flax.jax_utils.unreplicate(agent_state)
|
845 |
+
model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
|
846 |
+
with open(model_path, "wb") as f:
|
847 |
+
f.write(
|
848 |
+
flax.serialization.to_bytes(
|
849 |
+
[
|
850 |
+
vars(args),
|
851 |
+
[
|
852 |
+
agent_state.params.network_params,
|
853 |
+
agent_state.params.actor_params,
|
854 |
+
agent_state.params.critic_params,
|
855 |
+
],
|
856 |
+
]
|
857 |
+
)
|
858 |
+
)
|
859 |
+
print(f"model saved to {model_path}")
|
860 |
+
from cleanrl_utils.evals.ppo_envpool_jax_eval import evaluate
|
861 |
+
|
862 |
+
episodic_returns = evaluate(
|
863 |
+
model_path,
|
864 |
+
make_env,
|
865 |
+
args.env_id,
|
866 |
+
eval_episodes=10,
|
867 |
+
run_name=f"{run_name}-eval",
|
868 |
+
Model=(Network, Actor, Critic),
|
869 |
+
)
|
870 |
+
for idx, episodic_return in enumerate(episodic_returns):
|
871 |
+
writer.add_scalar("eval/episodic_return", episodic_return, idx)
|
872 |
+
|
873 |
+
if args.upload_model:
|
874 |
+
from cleanrl_utils.huggingface import push_to_hub
|
875 |
+
|
876 |
+
repo_name = f"{args.env_id}-{args.exp_name}-seed{args.seed}"
|
877 |
+
repo_id = f"{args.hf_entity}/{repo_name}" if args.hf_entity else repo_name
|
878 |
+
push_to_hub(
|
879 |
+
args,
|
880 |
+
episodic_returns,
|
881 |
+
repo_id,
|
882 |
+
"PPO",
|
883 |
+
f"runs/{run_name}",
|
884 |
+
f"videos/{run_name}-eval",
|
885 |
+
extra_dependencies=["jax", "envpool", "atari"],
|
886 |
+
)
|
887 |
+
|
888 |
+
envs.close()
|
889 |
+
writer.close()
|
videos/Pitfall-v5__sebulba_ppo_envpool__1__552a94c8-c755-4d9a-a517-2968c95f4ca4-eval/0.mp4
ADDED
Binary file (28.9 kB). View file
|
|