vwxyzjn commited on
Commit
ab1f628
1 Parent(s): 76cd0a6

pushing model

Browse files
.gitattributes CHANGED
@@ -32,3 +32,6 @@ 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
+ videos/Breakout-v5__ppo_atari_envpool_xla_jax_scan__2__1672591976-eval/0.mp4 filter=lfs diff=lfs merge=lfs -text
36
+ replay.mp4 filter=lfs diff=lfs merge=lfs -text
37
+ ppo_atari_envpool_xla_jax_scan.cleanrl_model filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - Breakout-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: Breakout-v5
16
+ type: Breakout-v5
17
+ metrics:
18
+ - type: mean_reward
19
+ value: 456.90 +/- 125.44
20
+ name: mean_reward
21
+ verified: false
22
+ ---
23
+
24
+ # (CleanRL) **PPO** Agent Playing **Breakout-v5**
25
+
26
+ This is a trained model of a PPO agent playing Breakout-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/ppo_atari_envpool_xla_jax_scan.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[ppo_atari_envpool_xla_jax_scan]"
36
+ python -m cleanrl_utils.enjoy --exp-name ppo_atari_envpool_xla_jax_scan --env-id Breakout-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/vwxyzjn/Breakout-v5-ppo_atari_envpool_xla_jax_scan-seed2/raw/main/ppo_atari_envpool_xla_jax_scan.py
46
+ curl -OL https://huggingface.co/vwxyzjn/Breakout-v5-ppo_atari_envpool_xla_jax_scan-seed2/raw/main/pyproject.toml
47
+ curl -OL https://huggingface.co/vwxyzjn/Breakout-v5-ppo_atari_envpool_xla_jax_scan-seed2/raw/main/poetry.lock
48
+ poetry install --all-extras
49
+ python ppo_atari_envpool_xla_jax_scan.py --track --save-model --upload-model --env-id Breakout-v5 --seed 2
50
+ ```
51
+
52
+ # Hyperparameters
53
+ ```python
54
+ {'anneal_lr': True,
55
+ 'batch_size': 1024,
56
+ 'capture_video': False,
57
+ 'clip_coef': 0.1,
58
+ 'cuda': True,
59
+ 'ent_coef': 0.01,
60
+ 'env_id': 'Breakout-v5',
61
+ 'exp_name': 'ppo_atari_envpool_xla_jax_scan',
62
+ 'gae_lambda': 0.95,
63
+ 'gamma': 0.99,
64
+ 'hf_entity': '',
65
+ 'learning_rate': 0.00025,
66
+ 'max_grad_norm': 0.5,
67
+ 'minibatch_size': 256,
68
+ 'norm_adv': True,
69
+ 'num_envs': 8,
70
+ 'num_minibatches': 4,
71
+ 'num_steps': 128,
72
+ 'num_updates': 9765,
73
+ 'save_model': True,
74
+ 'seed': 2,
75
+ 'target_kl': None,
76
+ 'torch_deterministic': True,
77
+ 'total_timesteps': 10000000,
78
+ 'track': True,
79
+ 'update_epochs': 4,
80
+ 'upload_model': True,
81
+ 'vf_coef': 0.5,
82
+ 'wandb_entity': None,
83
+ 'wandb_project_name': 'cleanRL'}
84
+ ```
85
+
events.out.tfevents.1672591978.pop-os.2768365.0 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c03da540173bc5550a066b2f5e59ea87c065df352890ad6b82fac2693e9c8536
3
+ size 5693196
poetry.lock ADDED
The diff for this file is too large to render. See raw diff
 
ppo_atari_envpool_xla_jax_scan.cleanrl_model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2815f9e945c931d5aefaa16fed286c713c43a9577e286dbe1c80dc8fd56a72fc
3
+ size 6747662
ppo_atari_envpool_xla_jax_scan.py ADDED
@@ -0,0 +1,514 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_atari_envpool_xla_jaxpy
2
+ import argparse
3
+ import os
4
+ import random
5
+ import time
6
+ from distutils.util import strtobool
7
+ from functools import partial
8
+ from typing import Sequence
9
+
10
+ os.environ[
11
+ "XLA_PYTHON_CLIENT_MEM_FRACTION"
12
+ ] = "0.7" # see https://github.com/google/jax/discussions/6332#discussioncomment-1279991
13
+
14
+ import envpool
15
+ import flax
16
+ import flax.linen as nn
17
+ import gym
18
+ import jax
19
+ import jax.numpy as jnp
20
+ import numpy as np
21
+ import optax
22
+ from flax.linen.initializers import constant, orthogonal
23
+ from flax.training.train_state import TrainState
24
+ from torch.utils.tensorboard import SummaryWriter
25
+
26
+
27
+ def parse_args():
28
+ # fmt: off
29
+ parser = argparse.ArgumentParser()
30
+ parser.add_argument("--exp-name", type=str, default=os.path.basename(__file__).rstrip(".py"),
31
+ help="the name of this experiment")
32
+ parser.add_argument("--seed", type=int, default=1,
33
+ help="seed of the experiment")
34
+ parser.add_argument("--torch-deterministic", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
35
+ help="if toggled, `torch.backends.cudnn.deterministic=False`")
36
+ parser.add_argument("--cuda", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
37
+ help="if toggled, cuda will be enabled by default")
38
+ parser.add_argument("--track", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
39
+ help="if toggled, this experiment will be tracked with Weights and Biases")
40
+ parser.add_argument("--wandb-project-name", type=str, default="cleanRL",
41
+ help="the wandb's project name")
42
+ parser.add_argument("--wandb-entity", type=str, default=None,
43
+ help="the entity (team) of wandb's project")
44
+ parser.add_argument("--capture-video", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
45
+ help="whether to capture videos of the agent performances (check out `videos` folder)")
46
+ parser.add_argument("--save-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
47
+ help="whether to save model into the `runs/{run_name}` folder")
48
+ parser.add_argument("--upload-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
49
+ help="whether to upload the saved model to huggingface")
50
+ parser.add_argument("--hf-entity", type=str, default="",
51
+ help="the user or org name of the model repository from the Hugging Face Hub")
52
+
53
+ # Algorithm specific arguments
54
+ parser.add_argument("--env-id", type=str, default="Pong-v5",
55
+ help="the id of the environment")
56
+ parser.add_argument("--total-timesteps", type=int, default=10000000,
57
+ help="total timesteps of the experiments")
58
+ parser.add_argument("--learning-rate", type=float, default=2.5e-4,
59
+ help="the learning rate of the optimizer")
60
+ parser.add_argument("--num-envs", type=int, default=8,
61
+ help="the number of parallel game environments")
62
+ parser.add_argument("--num-steps", type=int, default=128,
63
+ help="the number of steps to run in each environment per policy rollout")
64
+ parser.add_argument("--anneal-lr", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
65
+ help="Toggle learning rate annealing for policy and value networks")
66
+ parser.add_argument("--gamma", type=float, default=0.99,
67
+ help="the discount factor gamma")
68
+ parser.add_argument("--gae-lambda", type=float, default=0.95,
69
+ help="the lambda for the general advantage estimation")
70
+ parser.add_argument("--num-minibatches", type=int, default=4,
71
+ help="the number of mini-batches")
72
+ parser.add_argument("--update-epochs", type=int, default=4,
73
+ help="the K epochs to update the policy")
74
+ parser.add_argument("--norm-adv", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
75
+ help="Toggles advantages normalization")
76
+ parser.add_argument("--clip-coef", type=float, default=0.1,
77
+ help="the surrogate clipping coefficient")
78
+ parser.add_argument("--ent-coef", type=float, default=0.01,
79
+ help="coefficient of the entropy")
80
+ parser.add_argument("--vf-coef", type=float, default=0.5,
81
+ help="coefficient of the value function")
82
+ parser.add_argument("--max-grad-norm", type=float, default=0.5,
83
+ help="the maximum norm for the gradient clipping")
84
+ parser.add_argument("--target-kl", type=float, default=None,
85
+ help="the target KL divergence threshold")
86
+ args = parser.parse_args()
87
+ args.batch_size = int(args.num_envs * args.num_steps)
88
+ args.minibatch_size = int(args.batch_size // args.num_minibatches)
89
+ args.num_updates = args.total_timesteps // args.batch_size
90
+ # fmt: on
91
+ return args
92
+
93
+
94
+ def make_env(env_id, seed, num_envs):
95
+ def thunk():
96
+ envs = envpool.make(
97
+ env_id,
98
+ env_type="gym",
99
+ num_envs=num_envs,
100
+ episodic_life=True,
101
+ reward_clip=True,
102
+ seed=seed,
103
+ )
104
+ envs.num_envs = num_envs
105
+ envs.single_action_space = envs.action_space
106
+ envs.single_observation_space = envs.observation_space
107
+ envs.is_vector_env = True
108
+ return envs
109
+
110
+ return thunk
111
+
112
+
113
+ class Network(nn.Module):
114
+ @nn.compact
115
+ def __call__(self, x):
116
+ x = jnp.transpose(x, (0, 2, 3, 1))
117
+ x = x / (255.0)
118
+ x = nn.Conv(
119
+ 32,
120
+ kernel_size=(8, 8),
121
+ strides=(4, 4),
122
+ padding="VALID",
123
+ kernel_init=orthogonal(np.sqrt(2)),
124
+ bias_init=constant(0.0),
125
+ )(x)
126
+ x = nn.relu(x)
127
+ x = nn.Conv(
128
+ 64,
129
+ kernel_size=(4, 4),
130
+ strides=(2, 2),
131
+ padding="VALID",
132
+ kernel_init=orthogonal(np.sqrt(2)),
133
+ bias_init=constant(0.0),
134
+ )(x)
135
+ x = nn.relu(x)
136
+ x = nn.Conv(
137
+ 64,
138
+ kernel_size=(3, 3),
139
+ strides=(1, 1),
140
+ padding="VALID",
141
+ kernel_init=orthogonal(np.sqrt(2)),
142
+ bias_init=constant(0.0),
143
+ )(x)
144
+ x = nn.relu(x)
145
+ x = x.reshape((x.shape[0], -1))
146
+ x = nn.Dense(512, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(x)
147
+ x = nn.relu(x)
148
+ return x
149
+
150
+
151
+ class Critic(nn.Module):
152
+ @nn.compact
153
+ def __call__(self, x):
154
+ return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
155
+
156
+
157
+ class Actor(nn.Module):
158
+ action_dim: Sequence[int]
159
+
160
+ @nn.compact
161
+ def __call__(self, x):
162
+ return nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
163
+
164
+
165
+ @flax.struct.dataclass
166
+ class AgentParams:
167
+ network_params: flax.core.FrozenDict
168
+ actor_params: flax.core.FrozenDict
169
+ critic_params: flax.core.FrozenDict
170
+
171
+
172
+ @flax.struct.dataclass
173
+ class Storage:
174
+ obs: jnp.array
175
+ actions: jnp.array
176
+ logprobs: jnp.array
177
+ dones: jnp.array
178
+ values: jnp.array
179
+ advantages: jnp.array
180
+ returns: jnp.array
181
+ rewards: jnp.array
182
+
183
+
184
+ @flax.struct.dataclass
185
+ class EpisodeStatistics:
186
+ episode_returns: jnp.array
187
+ episode_lengths: jnp.array
188
+ returned_episode_returns: jnp.array
189
+ returned_episode_lengths: jnp.array
190
+
191
+
192
+ if __name__ == "__main__":
193
+ args = parse_args()
194
+ run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
195
+ if args.track:
196
+ import wandb
197
+
198
+ wandb.init(
199
+ project=args.wandb_project_name,
200
+ entity=args.wandb_entity,
201
+ sync_tensorboard=True,
202
+ config=vars(args),
203
+ name=run_name,
204
+ monitor_gym=True,
205
+ save_code=True,
206
+ )
207
+ writer = SummaryWriter(f"runs/{run_name}")
208
+ writer.add_text(
209
+ "hyperparameters",
210
+ "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
211
+ )
212
+
213
+ # TRY NOT TO MODIFY: seeding
214
+ random.seed(args.seed)
215
+ np.random.seed(args.seed)
216
+ key = jax.random.PRNGKey(args.seed)
217
+ key, network_key, actor_key, critic_key = jax.random.split(key, 4)
218
+
219
+ # env setup
220
+ envs = make_env(args.env_id, args.seed, args.num_envs)()
221
+ episode_stats = EpisodeStatistics(
222
+ episode_returns=jnp.zeros(args.num_envs, dtype=jnp.float32),
223
+ episode_lengths=jnp.zeros(args.num_envs, dtype=jnp.int32),
224
+ returned_episode_returns=jnp.zeros(args.num_envs, dtype=jnp.float32),
225
+ returned_episode_lengths=jnp.zeros(args.num_envs, dtype=jnp.int32),
226
+ )
227
+ handle, recv, send, step_env = envs.xla()
228
+
229
+ def step_env_wrappeed(episode_stats, handle, action):
230
+ handle, (next_obs, reward, next_done, info) = step_env(handle, action)
231
+ new_episode_return = episode_stats.episode_returns + info["reward"]
232
+ new_episode_length = episode_stats.episode_lengths + 1
233
+ episode_stats = episode_stats.replace(
234
+ episode_returns=(new_episode_return) * (1 - info["terminated"]) * (1 - info["TimeLimit.truncated"]),
235
+ episode_lengths=(new_episode_length) * (1 - info["terminated"]) * (1 - info["TimeLimit.truncated"]),
236
+ # only update the `returned_episode_returns` if the episode is done
237
+ returned_episode_returns=jnp.where(
238
+ info["terminated"] + info["TimeLimit.truncated"], new_episode_return, episode_stats.returned_episode_returns
239
+ ),
240
+ returned_episode_lengths=jnp.where(
241
+ info["terminated"] + info["TimeLimit.truncated"], new_episode_length, episode_stats.returned_episode_lengths
242
+ ),
243
+ )
244
+ return episode_stats, handle, (next_obs, reward, next_done, info)
245
+
246
+ assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
247
+
248
+ def linear_schedule(count):
249
+ # anneal learning rate linearly after one training iteration which contains
250
+ # (args.num_minibatches * args.update_epochs) gradient updates
251
+ frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_updates
252
+ return args.learning_rate * frac
253
+
254
+ network = Network()
255
+ actor = Actor(action_dim=envs.single_action_space.n)
256
+ critic = Critic()
257
+ network_params = network.init(network_key, np.array([envs.single_observation_space.sample()]))
258
+ agent_state = TrainState.create(
259
+ apply_fn=None,
260
+ params=AgentParams(
261
+ network_params,
262
+ actor.init(actor_key, network.apply(network_params, np.array([envs.single_observation_space.sample()]))),
263
+ critic.init(critic_key, network.apply(network_params, np.array([envs.single_observation_space.sample()]))),
264
+ ),
265
+ tx=optax.chain(
266
+ optax.clip_by_global_norm(args.max_grad_norm),
267
+ optax.inject_hyperparams(optax.adam)(
268
+ learning_rate=linear_schedule if args.anneal_lr else args.learning_rate, eps=1e-5
269
+ ),
270
+ ),
271
+ )
272
+ network.apply = jax.jit(network.apply)
273
+ actor.apply = jax.jit(actor.apply)
274
+ critic.apply = jax.jit(critic.apply)
275
+
276
+ @jax.jit
277
+ def get_action_and_value(
278
+ agent_state: TrainState,
279
+ next_obs: np.ndarray,
280
+ key: jax.random.PRNGKey,
281
+ ):
282
+ """sample action, calculate value, logprob, entropy, and update storage"""
283
+ hidden = network.apply(agent_state.params.network_params, next_obs)
284
+ logits = actor.apply(agent_state.params.actor_params, hidden)
285
+ # sample action: Gumbel-softmax trick
286
+ # see https://stats.stackexchange.com/questions/359442/sampling-from-a-categorical-distribution
287
+ key, subkey = jax.random.split(key)
288
+ u = jax.random.uniform(subkey, shape=logits.shape)
289
+ action = jnp.argmax(logits - jnp.log(-jnp.log(u)), axis=1)
290
+ logprob = jax.nn.log_softmax(logits)[jnp.arange(action.shape[0]), action]
291
+ value = critic.apply(agent_state.params.critic_params, hidden)
292
+ return action, logprob, value.squeeze(1), key
293
+
294
+ @jax.jit
295
+ def get_action_and_value2(
296
+ params: flax.core.FrozenDict,
297
+ x: np.ndarray,
298
+ action: np.ndarray,
299
+ ):
300
+ """calculate value, logprob of supplied `action`, and entropy"""
301
+ hidden = network.apply(params.network_params, x)
302
+ logits = actor.apply(params.actor_params, hidden)
303
+ logprob = jax.nn.log_softmax(logits)[jnp.arange(action.shape[0]), action]
304
+ # normalize the logits https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/
305
+ logits = logits - jax.scipy.special.logsumexp(logits, axis=-1, keepdims=True)
306
+ logits = logits.clip(min=jnp.finfo(logits.dtype).min)
307
+ p_log_p = logits * jax.nn.softmax(logits)
308
+ entropy = -p_log_p.sum(-1)
309
+ value = critic.apply(params.critic_params, hidden).squeeze()
310
+ return logprob, entropy, value
311
+
312
+ def compute_gae_once(carry, inp, gamma, gae_lambda):
313
+ advantages = carry
314
+ nextdone, nextvalues, curvalues, reward = inp
315
+ nextnonterminal = 1.0 - nextdone
316
+
317
+ delta = reward + gamma * nextvalues * nextnonterminal - curvalues
318
+ advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
319
+ return advantages, advantages
320
+
321
+ compute_gae_once = partial(compute_gae_once, gamma=args.gamma, gae_lambda=args.gae_lambda)
322
+
323
+ @jax.jit
324
+ def compute_gae(
325
+ agent_state: TrainState,
326
+ next_obs: np.ndarray,
327
+ next_done: np.ndarray,
328
+ storage: Storage,
329
+ ):
330
+ next_value = critic.apply(
331
+ agent_state.params.critic_params, network.apply(agent_state.params.network_params, next_obs)
332
+ ).squeeze()
333
+
334
+ advantages = jnp.zeros((args.num_envs,))
335
+ dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
336
+ values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
337
+ _, advantages = jax.lax.scan(
338
+ compute_gae_once, advantages, (dones[1:], values[1:], values[:-1], storage.rewards), reverse=True
339
+ )
340
+ storage = storage.replace(
341
+ advantages=advantages,
342
+ returns=advantages + storage.values,
343
+ )
344
+ return storage
345
+
346
+ def ppo_loss(params, x, a, logp, mb_advantages, mb_returns):
347
+ newlogprob, entropy, newvalue = get_action_and_value2(params, x, a)
348
+ logratio = newlogprob - logp
349
+ ratio = jnp.exp(logratio)
350
+ approx_kl = ((ratio - 1) - logratio).mean()
351
+
352
+ if args.norm_adv:
353
+ mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
354
+
355
+ # Policy loss
356
+ pg_loss1 = -mb_advantages * ratio
357
+ pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
358
+ pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
359
+
360
+ # Value loss
361
+ v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
362
+
363
+ entropy_loss = entropy.mean()
364
+ loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
365
+ return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
366
+
367
+ ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True)
368
+
369
+ @jax.jit
370
+ def update_ppo(
371
+ agent_state: TrainState,
372
+ storage: Storage,
373
+ key: jax.random.PRNGKey,
374
+ ):
375
+ def update_epoch(carry, unused_inp):
376
+ agent_state, key = carry
377
+ key, subkey = jax.random.split(key)
378
+
379
+ def flatten(x):
380
+ return x.reshape((-1,) + x.shape[2:])
381
+
382
+ # taken from: https://github.com/google/brax/blob/main/brax/training/agents/ppo/train.py
383
+ def convert_data(x: jnp.ndarray):
384
+ x = jax.random.permutation(subkey, x)
385
+ x = jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
386
+ return x
387
+
388
+ flatten_storage = jax.tree_map(flatten, storage)
389
+ shuffled_storage = jax.tree_map(convert_data, flatten_storage)
390
+
391
+ def update_minibatch(agent_state, minibatch):
392
+ (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
393
+ agent_state.params,
394
+ minibatch.obs,
395
+ minibatch.actions,
396
+ minibatch.logprobs,
397
+ minibatch.advantages,
398
+ minibatch.returns,
399
+ )
400
+ agent_state = agent_state.apply_gradients(grads=grads)
401
+ return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
402
+
403
+ agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
404
+ update_minibatch, agent_state, shuffled_storage
405
+ )
406
+ return (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
407
+
408
+ (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
409
+ update_epoch, (agent_state, key), (), length=args.update_epochs
410
+ )
411
+ return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
412
+
413
+ # TRY NOT TO MODIFY: start the game
414
+ global_step = 0
415
+ start_time = time.time()
416
+ next_obs = envs.reset()
417
+ next_done = jnp.zeros(args.num_envs, dtype=jax.numpy.bool_)
418
+
419
+ # based on https://github.dev/google/evojax/blob/0625d875262011d8e1b6aa32566b236f44b4da66/evojax/sim_mgr.py
420
+ def step_once(carry, step, env_step_fn):
421
+ agent_state, episode_stats, obs, done, key, handle = carry
422
+ action, logprob, value, key = get_action_and_value(agent_state, obs, key)
423
+
424
+ episode_stats, handle, (next_obs, reward, next_done, _) = env_step_fn(episode_stats, handle, action)
425
+ storage = Storage(
426
+ obs=obs,
427
+ actions=action,
428
+ logprobs=logprob,
429
+ dones=done,
430
+ values=value,
431
+ rewards=reward,
432
+ returns=jnp.zeros_like(reward),
433
+ advantages=jnp.zeros_like(reward),
434
+ )
435
+ return ((agent_state, episode_stats, next_obs, next_done, key, handle), storage)
436
+
437
+ def rollout(agent_state, episode_stats, next_obs, next_done, key, handle, step_once_fn, max_steps):
438
+ (agent_state, episode_stats, next_obs, next_done, key, handle), storage = jax.lax.scan(
439
+ step_once_fn, (agent_state, episode_stats, next_obs, next_done, key, handle), (), max_steps
440
+ )
441
+ return agent_state, episode_stats, next_obs, next_done, storage, key, handle
442
+
443
+ rollout = partial(rollout, step_once_fn=partial(step_once, env_step_fn=step_env_wrappeed), max_steps=args.num_steps)
444
+
445
+ for update in range(1, args.num_updates + 1):
446
+ update_time_start = time.time()
447
+ agent_state, episode_stats, next_obs, next_done, storage, key, handle = rollout(
448
+ agent_state, episode_stats, next_obs, next_done, key, handle
449
+ )
450
+ global_step += args.num_steps * args.num_envs
451
+ storage = compute_gae(agent_state, next_obs, next_done, storage)
452
+ agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key = update_ppo(
453
+ agent_state,
454
+ storage,
455
+ key,
456
+ )
457
+ avg_episodic_return = np.mean(jax.device_get(episode_stats.returned_episode_returns))
458
+ print(f"global_step={global_step}, avg_episodic_return={avg_episodic_return}")
459
+
460
+ # TRY NOT TO MODIFY: record rewards for plotting purposes
461
+ writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
462
+ writer.add_scalar(
463
+ "charts/avg_episodic_length", np.mean(jax.device_get(episode_stats.returned_episode_lengths)), global_step
464
+ )
465
+ writer.add_scalar("charts/learning_rate", agent_state.opt_state[1].hyperparams["learning_rate"].item(), global_step)
466
+ writer.add_scalar("losses/value_loss", v_loss[-1, -1].item(), global_step)
467
+ writer.add_scalar("losses/policy_loss", pg_loss[-1, -1].item(), global_step)
468
+ writer.add_scalar("losses/entropy", entropy_loss[-1, -1].item(), global_step)
469
+ writer.add_scalar("losses/approx_kl", approx_kl[-1, -1].item(), global_step)
470
+ writer.add_scalar("losses/loss", loss[-1, -1].item(), global_step)
471
+ print("SPS:", int(global_step / (time.time() - start_time)))
472
+ writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
473
+ writer.add_scalar(
474
+ "charts/SPS_update", int(args.num_envs * args.num_steps / (time.time() - update_time_start)), global_step
475
+ )
476
+
477
+ if args.save_model:
478
+ model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
479
+ with open(model_path, "wb") as f:
480
+ f.write(
481
+ flax.serialization.to_bytes(
482
+ [
483
+ vars(args),
484
+ [
485
+ agent_state.params.network_params,
486
+ agent_state.params.actor_params,
487
+ agent_state.params.critic_params,
488
+ ],
489
+ ]
490
+ )
491
+ )
492
+ print(f"model saved to {model_path}")
493
+ from cleanrl_utils.evals.ppo_envpool_jax_eval import evaluate
494
+
495
+ episodic_returns = evaluate(
496
+ model_path,
497
+ make_env,
498
+ args.env_id,
499
+ eval_episodes=10,
500
+ run_name=f"{run_name}-eval",
501
+ Model=(Network, Actor, Critic),
502
+ )
503
+ for idx, episodic_return in enumerate(episodic_returns):
504
+ writer.add_scalar("eval/episodic_return", episodic_return, idx)
505
+
506
+ if args.upload_model:
507
+ from cleanrl_utils.huggingface import push_to_hub
508
+
509
+ repo_name = f"{args.env_id}-{args.exp_name}-seed{args.seed}"
510
+ repo_id = f"{args.hf_entity}/{repo_name}" if args.hf_entity else repo_name
511
+ push_to_hub(args, episodic_returns, repo_id, "PPO", f"runs/{run_name}", f"videos/{run_name}-eval")
512
+
513
+ envs.close()
514
+ writer.close()
pyproject.toml ADDED
@@ -0,0 +1,178 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [tool.poetry]
2
+ name = "cleanrl-test"
3
+ version = "1.1.2"
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.7.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.7.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
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dec6bf236f0188975927c24f25eefe7c65b4763b01fd50ebd01c423604a19246
3
+ size 3974488
videos/Breakout-v5__ppo_atari_envpool_xla_jax_scan__2__1672591976-eval/0.mp4 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dec6bf236f0188975927c24f25eefe7c65b4763b01fd50ebd01c423604a19246
3
+ size 3974488