vwxyzjn commited on
Commit
1b96702
1 Parent(s): 106eded

pushing model

Browse files
.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
+ - Surround-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: Surround-v5
16
+ type: Surround-v5
17
+ metrics:
18
+ - type: mean_reward
19
+ value: 6.60 +/- 1.56
20
+ name: mean_reward
21
+ verified: false
22
+ ---
23
+
24
+ # (CleanRL) **PPO** Agent Playing **Surround-v5**
25
+
26
+ This is a trained model of a PPO agent playing Surround-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 Surround-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/Surround-v5-sebulba_ppo_envpool-seed1/raw/main/sebulba_ppo_envpool.py
46
+ curl -OL https://huggingface.co/cleanrl/Surround-v5-sebulba_ppo_envpool-seed1/raw/main/pyproject.toml
47
+ curl -OL https://huggingface.co/cleanrl/Surround-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 Surround-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': 'Surround-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.1675616260.ip-26-0-141-241.895734.0 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:96d19c5717b4f08b77077fcb7ad0f17ea5a91279ed5793f96929b8924a614934
3
+ size 9358615
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 <[email protected]>"]
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 (340 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:4245dc608c5e46ca35e3698661b1cc8af3c2783448a1bb078261dc3ea9f23e01
3
+ size 4378342
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/Surround-v5__sebulba_ppo_envpool__1__b29bc100-0528-4d6b-bee8-8a6f2cd8f3ab-eval/0.mp4 ADDED
Binary file (340 kB). View file