jacksonhack commited on
Commit
608f422
1 Parent(s): 205b8a2

pushing model

Browse files
README.md ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - CartPole-v1
4
+ - deep-reinforcement-learning
5
+ - reinforcement-learning
6
+ - custom-implementation
7
+ library_name: cleanrl
8
+ model-index:
9
+ - name: DQN
10
+ results:
11
+ - task:
12
+ type: reinforcement-learning
13
+ name: reinforcement-learning
14
+ dataset:
15
+ name: CartPole-v1
16
+ type: CartPole-v1
17
+ metrics:
18
+ - type: mean_reward
19
+ value: 88.80 +/- 47.92
20
+ name: mean_reward
21
+ verified: false
22
+ ---
23
+
24
+ # (CleanRL) **DQN** Agent Playing **CartPole-v1**
25
+
26
+ This is a trained model of a DQN agent playing CartPole-v1.
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/dqn.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[dqn]"
36
+ python -m cleanrl_utils.enjoy --exp-name dqn --env-id CartPole-v1
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/jacksonhack/CartPole-v1-dqn-seed1/raw/main/dqn.py
46
+ curl -OL https://huggingface.co/jacksonhack/CartPole-v1-dqn-seed1/raw/main/pyproject.toml
47
+ curl -OL https://huggingface.co/jacksonhack/CartPole-v1-dqn-seed1/raw/main/poetry.lock
48
+ poetry install --all-extras
49
+ python dqn.py --save-model --upload-model --total_timesteps 1000
50
+ ```
51
+
52
+ # Hyperparameters
53
+ ```python
54
+ {'batch_size': 128,
55
+ 'buffer_size': 10000,
56
+ 'capture_video': False,
57
+ 'cuda': True,
58
+ 'end_e': 0.05,
59
+ 'env_id': 'CartPole-v1',
60
+ 'exp_name': 'dqn',
61
+ 'exploration_fraction': 0.5,
62
+ 'gamma': 0.99,
63
+ 'hf_entity': 'jacksonhack',
64
+ 'learning_rate': 0.00025,
65
+ 'learning_starts': 10000,
66
+ 'num_envs': 1,
67
+ 'save_model': True,
68
+ 'seed': 1,
69
+ 'start_e': 1,
70
+ 'target_network_frequency': 500,
71
+ 'tau': 1.0,
72
+ 'torch_deterministic': True,
73
+ 'total_timesteps': 1000,
74
+ 'track': False,
75
+ 'train_frequency': 10,
76
+ 'upload_model': True,
77
+ 'wandb_entity': None,
78
+ 'wandb_project_name': 'cleanRL'}
79
+ ```
80
+
dqn.cleanrl_model ADDED
Binary file (46.2 kB). View file
 
dqn.py ADDED
@@ -0,0 +1,276 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ *Filename :dqn.py
3
+ *Description :
4
+ *Time :2024/11/13 18:34:33
5
+ *Author :jackson
6
+ *Version :1.0
7
+ """
8
+
9
+ import os
10
+ import random
11
+ import time
12
+ from dataclasses import dataclass
13
+
14
+ import gymnasium as gym
15
+ import numpy as np
16
+ import torch
17
+ import torch.nn as nn
18
+ import torch.nn.functional as F
19
+ import torch.optim as optim
20
+ import tyro
21
+ from stable_baselines3.common.buffers import ReplayBuffer
22
+
23
+ from torch.utils.tensorboard import SummaryWriter
24
+
25
+
26
+ @dataclass
27
+ class Args:
28
+ exp_name: str = os.path.basename(__file__)[: -len(".py")]
29
+ """the name of this experiment"""
30
+ seed: int = 1
31
+ """seed of the experiment"""
32
+ torch_deterministic: bool = True
33
+ """if toggled, `torch.backends.cudnn.deterministic=False`"""
34
+ cuda: bool = True
35
+ """if toggled, cuda will be enabled by default"""
36
+ track: bool = False
37
+ """if toggled, this experiment will be tracked with Weights and Biases"""
38
+ wandb_project_name: str = "cleanRL"
39
+ """the wandb's project name"""
40
+ wandb_entity: str = None
41
+ """the entity (team) of wandb's project"""
42
+ capture_video: bool = False
43
+ """whether to capture videos of the agent performances (check out `videos` folder)"""
44
+ save_model: bool = False
45
+ """whether to save model into the `runs/{run_name}` folder"""
46
+ upload_model: bool = False
47
+ """whether to upload the saved model to huggingface"""
48
+ hf_entity: str = "jacksonhack"
49
+ """the user or org name of the model repository from the Hugging Face Hub"""
50
+
51
+ # Algorithm specific arguments
52
+ env_id: str = "CartPole-v1"
53
+ """the id of the environment"""
54
+ total_timesteps: int = 500000
55
+ """total timesteps of the experiments"""
56
+ learning_rate: float = 2.5e-4
57
+ """the learning rate of the optimizer"""
58
+ num_envs: int = 1
59
+ """the number of parallel game environments"""
60
+ buffer_size: int = 10000
61
+ """the replay memory buffer size"""
62
+ gamma: float = 0.99
63
+ """the discount factor gamma"""
64
+ tau: float = 1.0
65
+ """the target network update rate"""
66
+ target_network_frequency: int = 500
67
+ """the timesteps it takes to update the target network"""
68
+ batch_size: int = 128
69
+ """the batch size of sample from the reply memory"""
70
+ start_e: float = 1
71
+ """the starting epsilon for exploration"""
72
+ end_e: float = 0.05
73
+ """the ending epsilon for exploration"""
74
+ exploration_fraction: float = 0.5
75
+ """the fraction of `total-timesteps` it takes from start-e to go end-e"""
76
+ learning_starts: int = 10000
77
+ """timestep to start learning"""
78
+ train_frequency: int = 10
79
+ """the frequency of training"""
80
+
81
+
82
+ def make_env(env_id, seed, idx, capture_video, run_name):
83
+ def thunk():
84
+ if capture_video and idx == 0:
85
+ env = gym.make(env_id, render_mode="rgb_array")
86
+ env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
87
+ else:
88
+ env = gym.make(env_id)
89
+ env = gym.wrappers.RecordEpisodeStatistics(env)
90
+ env.action_space.seed(seed)
91
+
92
+ return env
93
+
94
+ return thunk
95
+
96
+
97
+ class QNetwork(nn.Module):
98
+ def __init__(self, env):
99
+ super().__init__()
100
+ self.network = nn.Sequential(
101
+ nn.Linear(np.array(env.single_observation_space.shape).prod(), 120),
102
+ nn.ReLU(),
103
+ nn.Linear(120, 84),
104
+ nn.ReLU(),
105
+ nn.Linear(84, env.single_action_space.n),
106
+ )
107
+
108
+ def forward(self, x):
109
+ return self.network(x)
110
+
111
+
112
+ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
113
+ slope = (end_e - start_e) / duration
114
+ return max(slope * t + start_e, end_e)
115
+
116
+
117
+ if __name__ == "__main__":
118
+ import stable_baselines3 as sb3
119
+
120
+ if sb3.__version__ < "2.0":
121
+ raise ValueError(
122
+ """Ongoing migration: run the following command to install the new dependencies:
123
+
124
+ poetry run pip install "stable_baselines3==2.0.0a1"
125
+ """
126
+ )
127
+ args = tyro.cli(Args)
128
+
129
+ assert args.num_envs == 1, "vectorized envs are not supported at the moment"
130
+
131
+ run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
132
+
133
+ if args.track:
134
+ import wandb
135
+
136
+ wandb.init(
137
+ project=args.wandb_project_name,
138
+ entity=args.wandb_entity,
139
+ sync_tensorboard=True,
140
+ config=vars(args),
141
+ name=run_name,
142
+ monitor_gym=True,
143
+ save_code=True,
144
+ )
145
+
146
+ writer = SummaryWriter(f"runs/{run_name}")
147
+ writer.add_text(
148
+ "hyperparameters",
149
+ "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
150
+ )
151
+
152
+ # TRY NOT TO MODIFY: seeding
153
+ random.seed(args.seed)
154
+ np.random.seed(args.seed)
155
+ torch.manual_seed(args.seed)
156
+ torch.backends.cudnn.deterministic = args.torch_deterministic
157
+
158
+ device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
159
+
160
+ # env setup
161
+
162
+ envs = gym.vector.SyncVectorEnv(
163
+ [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
164
+ )
165
+
166
+ assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
167
+
168
+ q_network = QNetwork(envs).to(device)
169
+ optimizer = optim.Adam(q_network.parameters(), lr=args.learning_rate)
170
+ target_network = QNetwork(envs).to(device)
171
+ target_network.load_state_dict(q_network.state_dict())
172
+
173
+ rb = ReplayBuffer(
174
+ args.buffer_size,
175
+ envs.single_observation_space,
176
+ envs.single_action_space,
177
+ device,
178
+ handle_timeout_termination=False,
179
+ )
180
+
181
+ start_time = time.time()
182
+
183
+ obs, _ = envs.reset(seed=args.seed)
184
+
185
+ for global_step in range(args.total_timesteps):
186
+ # ALGO LOGIC: put action logic here
187
+ epsilon = linear_schedule(
188
+ args.start_e, args.end_e, args.exploration_fraction * args.total_timesteps, global_step
189
+ )
190
+
191
+ if random.random() < epsilon:
192
+ actions = np.array([envs.single_action_space.sample() for _ in range(envs.num_envs)])
193
+ else:
194
+ q_values = q_network(torch.Tensor(obs).to(device))
195
+ actions = torch.argmax(q_values, dim=1).cpu().numpy()
196
+
197
+ # TRY NOT TO MODIFY: execute the game and log data.
198
+ next_obs, rewards, terminations, truncations, infos = envs.step(actions)
199
+
200
+ # TRY NOT TO MODIFY: record rewards for plotting purposes
201
+ if "final_info" in infos:
202
+ for info in infos["final_info"]:
203
+ if info and "episode" in info:
204
+ print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
205
+ writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
206
+ writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
207
+
208
+ # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
209
+ # 向量化环境会自己重置环境
210
+ real_next_obs = next_obs.copy()
211
+ for idx, trunc in enumerate(truncations):
212
+ if trunc:
213
+ # 将截断状态变为真实状态,确保算法获得更准确的信息
214
+ real_next_obs[idx] = infos["final_observation"][idx]
215
+
216
+ rb.add(obs, real_next_obs, actions, rewards, terminations, infos)
217
+
218
+ # TRY NOT TO MODIFY: CRUCIAL step easy to overlook
219
+ obs = next_obs
220
+
221
+ # ALGO LOGIC: training.
222
+
223
+ if global_step > args.learning_starts:
224
+ if global_step % args.train_frequency == 0:
225
+ data = rb.sample(args.batch_size)
226
+ with torch.no_grad():
227
+ target_max, _ = target_network(data.next_observations).max(dim=1) # tensor.max() 返回最大值及其索引
228
+ td_target = data.rewards.flatten() + args.gamma * target_max * (1 - data.dones.flatten())
229
+
230
+ old_val = q_network(data.observations).gather(1, data.actions).squeeze()
231
+ loss = F.mse_loss(td_target, old_val)
232
+
233
+ if global_step % 100 == 0:
234
+ writer.add_scalar("losses/td_loss", loss, global_step)
235
+ writer.add_scalar("losses/q_values", old_val.mean().item(), global_step)
236
+ # SPS: Step per second
237
+ print("SPS:", int(global_step / (time.time() - start_time)))
238
+ writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
239
+
240
+ # optimize the model
241
+ optimizer.zero_grad()
242
+ loss.backward()
243
+ optimizer.step()
244
+
245
+ # update target network
246
+ if global_step % args.target_network_frequency == 0:
247
+ for target_network_param, q_network_param in zip(target_network.parameters(), q_network.parameters()):
248
+ target_network_param.data.copy_(
249
+ args.tau * q_network_param.data + (1.0 - args.tau) * target_network_param.data
250
+ )
251
+
252
+ if args.save_model:
253
+ model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
254
+ torch.save(q_network.state_dict(), model_path)
255
+ print(f"model saved to {model_path}")
256
+ from utils.evals.dqn_eval import evaluate
257
+
258
+ episodic_returns = evaluate(
259
+ model_path,
260
+ make_env,
261
+ args.env_id,
262
+ eval_episodes=10,
263
+ run_name=f"{run_name}-eval",
264
+ Model=QNetwork,
265
+ device=device,
266
+ epsilon=0.05,
267
+ )
268
+ for idx, episodic_return in enumerate(episodic_returns):
269
+ writer.add_scalar("eval/episodic_return", episodic_return, idx)
270
+
271
+ if args.upload_model:
272
+ from utils.huggingface import push_to_hub
273
+
274
+ repo_name = f"{args.env_id}-{args.exp_name}-seed{args.seed}"
275
+ repo_id = f"{args.hf_entity}/{repo_name}" if args.hf_entity else repo_name
276
+ push_to_hub(args, episodic_returns, repo_id, "DQN", f"runs/{run_name}", f"videos/{run_name}-eval")
events.out.tfevents.1731593916.DESKTOP-3BC7099.139401.0 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d569722c75b6aee5c66c4e2901b35f4b26b8e2c8dcf935943a41e9f7fd612760
3
+ size 3435
replay.mp4 ADDED
Binary file (7.16 kB). View file
 
videos/CartPole-v1__dqn__1__1731593916-eval/rl-video-episode-0.mp4 ADDED
Binary file (7.37 kB). View file
 
videos/CartPole-v1__dqn__1__1731593916-eval/rl-video-episode-1.mp4 ADDED
Binary file (7.96 kB). View file
 
videos/CartPole-v1__dqn__1__1731593916-eval/rl-video-episode-8.mp4 ADDED
Binary file (7.16 kB). View file