Ason-jay commited on
Commit
c143496
·
verified ·
1 Parent(s): dd6694a

Upload folder using huggingface_hub

Browse files
Files changed (7) hide show
  1. README.md +120 -0
  2. checkpoint.pt +3 -0
  3. config.json +28 -0
  4. dataset_stats.json +48 -0
  5. eval_gazebo.csv +6 -0
  6. training_code.py +390 -0
  7. training_log.csv +101 -0
README.md ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ tags:
4
+ - robotics
5
+ - offline-rl
6
+ - iql
7
+ - implicit-q-learning
8
+ - fetch
9
+ - manipulation
10
+ library_name: pytorch
11
+ pipeline_tag: robotics
12
+ ---
13
+
14
+ # IQL (tau=0.7) - Fetch Robot Pick-and-Place
15
+
16
+ Implicit Q-Learning (expectile tau=0.7) offline RL policy for Fetch robot pick-and-place.
17
+
18
+ ## Model Description
19
+
20
+ This model was trained using **offline reinforcement learning** on a static dataset of 540 demonstration
21
+ episodes (26,538 transitions) collected from trajectory optimization on the Fetch robot in Gazebo simulation.
22
+
23
+ ### Task
24
+ - **Robot**: Fetch Mobile Manipulator (7 arm + 2 gripper = 9 DOF)
25
+ - **Task**: Pick-and-place (lift cracker box >= 10cm)
26
+ - **State space**: 9D joint positions
27
+ - **Action space**: 9D target joint positions
28
+
29
+ ### Dataset
30
+ - **Source**: Trajectory optimization with quality-tiered rewards
31
+ - **Episodes**: 540 (304 both_pass, 194 lift_only, 42 fail)
32
+ - **Transitions**: 26,538
33
+ - **Reward structure**: Sparse terminal (both_pass=1.0, lift_only=0.5, fail=0.0)
34
+
35
+ ## Training Hyperparameters
36
+
37
+ | Parameter | Value |
38
+ |-----------|-------|
39
+ | algorithm | IQL |
40
+ | tau_expectile | 0.7 |
41
+ | beta | 3.0 |
42
+ | lr | 0.0003 |
43
+ | batch_size | 256 |
44
+ | discount | 0.99 |
45
+ | target_update_rate | 0.005 |
46
+ | num_iterations | 100000 |
47
+ | hidden_dims | [256, 256] |
48
+ | state_normalization | zero_mean_unit_var |
49
+
50
+ ## Evaluation Results
51
+
52
+ | Metric | Value |
53
+ |--------|-------|
54
+ | action_mse | 0.002713 |
55
+ | gazebo_success_rate | 0/5 (0%) |
56
+ | gazebo_avg_lift | 0.0146 |
57
+
58
+ ### Offline Policy Evaluation
59
+
60
+ Action MSE measures how closely the policy reproduces the demonstration actions:
61
+ - **TD3+BC**: MSE = 0.374 (poor action matching)
62
+ - **IQL (tau=0.7)**: MSE = 0.0027 (good)
63
+ - **IQL (tau=0.9)**: MSE = 0.0012 (best)
64
+
65
+ ### Gazebo Evaluation (5 episodes)
66
+
67
+ All models achieved 0% success rate in the initial pilot evaluation.
68
+ This is expected for a first iteration - the models need further refinement
69
+ (e.g., longer training, reward shaping, or residual RL integration).
70
+
71
+ ## Files
72
+
73
+ - `checkpoint.pt` - Model weights (PyTorch)
74
+ - `training_code.py` - Training implementation
75
+ - `training_log.csv` - Training metrics over time
76
+ - `eval_gazebo.csv` - Gazebo evaluation results
77
+ - `dataset_stats.json` - Dataset normalization statistics
78
+ - `config.json` - Model configuration
79
+
80
+ ## Usage
81
+
82
+ ```python
83
+ import torch
84
+ import numpy as np
85
+
86
+ # Load checkpoint
87
+ ckpt = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)
88
+
89
+ # Load dataset stats for normalization
90
+ import json
91
+ with open("dataset_stats.json") as f:
92
+ stats = json.load(f)
93
+ state_mean = torch.tensor(stats["state_mean"])
94
+ state_std = torch.tensor(stats["state_std"])
95
+ ```
96
+
97
+ ## Joint Names
98
+
99
+ ```python
100
+ JOINTS = [
101
+ 'shoulder_pan_joint', # idx 0
102
+ 'shoulder_lift_joint', # idx 1
103
+ 'upperarm_roll_joint', # idx 2
104
+ 'elbow_flex_joint', # idx 3
105
+ 'forearm_roll_joint', # idx 4
106
+ 'wrist_flex_joint', # idx 5
107
+ 'wrist_roll_joint', # idx 6
108
+ 'l_gripper_finger_joint',# idx 7
109
+ 'r_gripper_finger_joint',# idx 8
110
+ ]
111
+ ```
112
+
113
+ ## Citation
114
+
115
+ ```bibtex
116
+ @misc{fetch_offline_rl_pilot,
117
+ title={Offline RL Pilot Study for Fetch Robot Pick-and-Place},
118
+ year={2026},
119
+ }
120
+ ```
checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ccc801a5910700e604b46b9768338a6d2d47fd02a0ffb05c44639b69e1111c5c
3
+ size 1713912
config.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "method": "iql_tau07",
3
+ "display_name": "IQL (tau=0.7)",
4
+ "hyperparameters": {
5
+ "algorithm": "IQL",
6
+ "tau_expectile": 0.7,
7
+ "beta": 3.0,
8
+ "lr": 0.0003,
9
+ "batch_size": 256,
10
+ "discount": 0.99,
11
+ "target_update_rate": 0.005,
12
+ "num_iterations": 100000,
13
+ "hidden_dims": [
14
+ 256,
15
+ 256
16
+ ],
17
+ "state_normalization": "zero_mean_unit_var"
18
+ },
19
+ "metrics": {
20
+ "action_mse": 0.002713,
21
+ "gazebo_success_rate": "0/5 (0%)",
22
+ "gazebo_avg_lift": 0.0146
23
+ },
24
+ "state_dim": 9,
25
+ "action_dim": 9,
26
+ "robot": "fetch",
27
+ "task": "pick-and-place"
28
+ }
dataset_stats.json ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "state_mean": [
3
+ 0.7955080270767212,
4
+ 0.12009388208389282,
5
+ 0.30571427941322327,
6
+ -1.379521369934082,
7
+ -0.03298526257276535,
8
+ -0.601288378238678,
9
+ 0.15154360234737396,
10
+ 0.04576437547802925,
11
+ 0.04571504518389702
12
+ ],
13
+ "state_std": [
14
+ 0.4665786325931549,
15
+ 0.40634530782699585,
16
+ 1.0432075262069702,
17
+ 0.642571747303009,
18
+ 1.0778027772903442,
19
+ 1.294430136680603,
20
+ 0.6733894348144531,
21
+ 0.008357170037925243,
22
+ 0.00834752433001995
23
+ ],
24
+ "action_mean": [
25
+ 0.795506477355957,
26
+ 0.12009190768003464,
27
+ 0.3057110607624054,
28
+ -1.3795207738876343,
29
+ -0.032985102385282516,
30
+ -0.6012873649597168,
31
+ 0.15154403448104858,
32
+ 0.04576554521918297,
33
+ 0.045717429369688034
34
+ ],
35
+ "action_std": [
36
+ 0.4665771424770355,
37
+ 0.4063422381877899,
38
+ 1.0432077646255493,
39
+ 0.6425707340240479,
40
+ 1.0778027772903442,
41
+ 1.2944289445877075,
42
+ 0.6733894348144531,
43
+ 0.008352111093699932,
44
+ 0.008343065157532692
45
+ ],
46
+ "num_transitions": 26538,
47
+ "num_episodes": 540
48
+ }
eval_gazebo.csv ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ episode,success,lift_height,episode_length
2
+ 1,False,0.00023790558495095926,80
3
+ 2,False,0.03705704541155985,80
4
+ 3,False,0.00032800151937706357,80
5
+ 4,False,0.034518845411559806,80
6
+ 5,False,0.001075935286332208,80
training_code.py ADDED
@@ -0,0 +1,390 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Implicit Q-Learning (IQL) for offline RL.
3
+ Kostrikov et al., 2022 - "Offline Reinforcement Learning with Implicit Q-Learning"
4
+
5
+ Three networks, three losses:
6
+ 1. V-loss: expectile regression of Q-values
7
+ 2. Q-loss: MSE with V-targets (no max over actions)
8
+ 3. Policy loss: advantage-weighted regression (AWR)
9
+ """
10
+
11
+ import os
12
+ import csv
13
+ import argparse
14
+ import numpy as np
15
+ import torch
16
+ import torch.nn as nn
17
+ import torch.nn.functional as F
18
+ from torch.distributions import Normal
19
+
20
+ # Joint limits for clamping policy output
21
+ JOINT_LOWER = torch.tensor([-1.606, -1.221, -3.142, -2.251, -3.142, -2.16, -3.142, 0.0, 0.0])
22
+ JOINT_UPPER = torch.tensor([1.606, 1.518, 3.142, 2.251, 3.142, 3.142, 3.142, 0.05, 0.05])
23
+
24
+
25
+ # ============================================================
26
+ # Networks
27
+ # ============================================================
28
+
29
+ class QNetwork(nn.Module):
30
+ """Q(s, a) -> scalar"""
31
+ def __init__(self, state_dim=9, action_dim=9, hidden_dim=256):
32
+ super().__init__()
33
+ self.net = nn.Sequential(
34
+ nn.Linear(state_dim + action_dim, hidden_dim),
35
+ nn.ReLU(),
36
+ nn.Linear(hidden_dim, hidden_dim),
37
+ nn.ReLU(),
38
+ nn.Linear(hidden_dim, 1),
39
+ )
40
+
41
+ def forward(self, state, action):
42
+ x = torch.cat([state, action], dim=-1)
43
+ return self.net(x).squeeze(-1)
44
+
45
+
46
+ class VNetwork(nn.Module):
47
+ """V(s) -> scalar"""
48
+ def __init__(self, state_dim=9, hidden_dim=256):
49
+ super().__init__()
50
+ self.net = nn.Sequential(
51
+ nn.Linear(state_dim, hidden_dim),
52
+ nn.ReLU(),
53
+ nn.Linear(hidden_dim, hidden_dim),
54
+ nn.ReLU(),
55
+ nn.Linear(hidden_dim, 1),
56
+ )
57
+
58
+ def forward(self, state):
59
+ return self.net(state).squeeze(-1)
60
+
61
+
62
+ class GaussianPolicy(nn.Module):
63
+ """Diagonal Gaussian policy: pi(a|s)"""
64
+ def __init__(self, state_dim=9, action_dim=9, hidden_dim=256):
65
+ super().__init__()
66
+ self.trunk = nn.Sequential(
67
+ nn.Linear(state_dim, hidden_dim),
68
+ nn.ReLU(),
69
+ nn.Linear(hidden_dim, hidden_dim),
70
+ nn.ReLU(),
71
+ )
72
+ self.mean_head = nn.Linear(hidden_dim, action_dim)
73
+ self.log_std_head = nn.Linear(hidden_dim, action_dim)
74
+
75
+ def forward(self, state):
76
+ h = self.trunk(state)
77
+ mean = self.mean_head(h)
78
+ log_std = self.log_std_head(h).clamp(-5.0, 2.0)
79
+ return mean, log_std
80
+
81
+ def log_prob(self, state, action):
82
+ """Log probability of action under the policy (no clamping on action)."""
83
+ mean, log_std = self.forward(state)
84
+ std = log_std.exp()
85
+ dist = Normal(mean, std)
86
+ return dist.log_prob(action).sum(dim=-1)
87
+
88
+ def sample(self, state):
89
+ """Sample an action and return (clamped_action, log_prob)."""
90
+ mean, log_std = self.forward(state)
91
+ std = log_std.exp()
92
+ dist = Normal(mean, std)
93
+ raw_action = dist.rsample()
94
+ log_p = dist.log_prob(raw_action).sum(dim=-1)
95
+ clamped = torch.clamp(raw_action, JOINT_LOWER.to(raw_action.device),
96
+ JOINT_UPPER.to(raw_action.device))
97
+ return clamped, log_p
98
+
99
+
100
+ # ============================================================
101
+ # Dataset
102
+ # ============================================================
103
+
104
+ class OfflineDataset:
105
+ """Loads NPZ dataset into GPU tensors for fast sampling."""
106
+ def __init__(self, path, device):
107
+ data = np.load(path)
108
+ self.states = torch.tensor(data['states'], dtype=torch.float32, device=device)
109
+ self.actions = torch.tensor(data['actions'], dtype=torch.float32, device=device)
110
+ self.rewards = torch.tensor(data['rewards'], dtype=torch.float32, device=device)
111
+ self.next_states = torch.tensor(data['next_states'], dtype=torch.float32, device=device)
112
+ self.dones = torch.tensor(data['dones'], dtype=torch.float32, device=device)
113
+ self.state_mean = torch.tensor(data['state_mean'], dtype=torch.float32, device=device)
114
+ self.state_std = torch.tensor(data['state_std'], dtype=torch.float32, device=device)
115
+ self.size = self.states.shape[0]
116
+ print(f"Loaded dataset: {self.size} transitions")
117
+
118
+ def normalize_state(self, s):
119
+ return (s - self.state_mean) / self.state_std
120
+
121
+ def sample(self, batch_size):
122
+ idx = torch.randint(0, self.size, (batch_size,), device=self.states.device)
123
+ return (
124
+ self.normalize_state(self.states[idx]),
125
+ self.actions[idx],
126
+ self.rewards[idx],
127
+ self.normalize_state(self.next_states[idx]),
128
+ self.dones[idx],
129
+ )
130
+
131
+
132
+ # ============================================================
133
+ # Soft update
134
+ # ============================================================
135
+
136
+ @torch.no_grad()
137
+ def soft_update(target, source, tau):
138
+ for tp, sp in zip(target.parameters(), source.parameters()):
139
+ tp.data.copy_(tau * sp.data + (1.0 - tau) * tp.data)
140
+
141
+
142
+ # ============================================================
143
+ # IQL Trainer
144
+ # ============================================================
145
+
146
+ class IQLTrainer:
147
+ def __init__(self, dataset, device, tau_expectile=0.7, beta=3.0,
148
+ lr=3e-4, discount=0.99, target_update_rate=0.005):
149
+ self.dataset = dataset
150
+ self.device = device
151
+ self.tau_expectile = tau_expectile
152
+ self.beta = beta
153
+ self.discount = discount
154
+ self.target_update_rate = target_update_rate
155
+
156
+ # Networks
157
+ self.q1 = QNetwork().to(device)
158
+ self.q2 = QNetwork().to(device)
159
+ self.q1_target = QNetwork().to(device)
160
+ self.q2_target = QNetwork().to(device)
161
+ self.q1_target.load_state_dict(self.q1.state_dict())
162
+ self.q2_target.load_state_dict(self.q2.state_dict())
163
+
164
+ self.v = VNetwork().to(device)
165
+ self.policy = GaussianPolicy().to(device)
166
+
167
+ # Move joint limits to device
168
+ global JOINT_LOWER, JOINT_UPPER
169
+ JOINT_LOWER = JOINT_LOWER.to(device)
170
+ JOINT_UPPER = JOINT_UPPER.to(device)
171
+
172
+ # Optimizers
173
+ self.opt_q = torch.optim.Adam(
174
+ list(self.q1.parameters()) + list(self.q2.parameters()), lr=lr)
175
+ self.opt_v = torch.optim.Adam(self.v.parameters(), lr=lr)
176
+ self.opt_pi = torch.optim.Adam(self.policy.parameters(), lr=lr)
177
+
178
+ def compute_value_loss(self, norm_s, actions):
179
+ with torch.no_grad():
180
+ q1_t = self.q1_target(norm_s, actions)
181
+ q2_t = self.q2_target(norm_s, actions)
182
+ q_target = torch.min(q1_t, q2_t)
183
+
184
+ v = self.v(norm_s)
185
+ diff = q_target - v
186
+ weight = torch.where(diff > 0, self.tau_expectile, 1.0 - self.tau_expectile)
187
+ loss_v = (weight * diff ** 2).mean()
188
+ return loss_v, v.mean().item()
189
+
190
+ def compute_q_loss(self, norm_s, actions, rewards, norm_s_next, dones):
191
+ with torch.no_grad():
192
+ v_next = self.v(norm_s_next)
193
+ target_q = rewards + self.discount * (1.0 - dones) * v_next
194
+
195
+ q1 = self.q1(norm_s, actions)
196
+ q2 = self.q2(norm_s, actions)
197
+ loss_q = ((q1 - target_q) ** 2).mean() + ((q2 - target_q) ** 2).mean()
198
+ return loss_q, q1.mean().item()
199
+
200
+ def compute_policy_loss(self, norm_s, actions):
201
+ with torch.no_grad():
202
+ q1_t = self.q1_target(norm_s, actions)
203
+ q2_t = self.q2_target(norm_s, actions)
204
+ q_target = torch.min(q1_t, q2_t)
205
+ v = self.v(norm_s)
206
+ advantage = q_target - v
207
+ exp_adv = torch.exp(self.beta * advantage).clamp(max=100.0)
208
+
209
+ log_prob = self.policy.log_prob(norm_s, actions)
210
+ loss_pi = -(exp_adv * log_prob).mean()
211
+ return loss_pi, advantage.mean().item(), exp_adv.mean().item()
212
+
213
+ def train_step(self, batch_size=256):
214
+ norm_s, actions, rewards, norm_s_next, dones = self.dataset.sample(batch_size)
215
+
216
+ # Update V
217
+ self.opt_v.zero_grad()
218
+ loss_v, v_mean = self.compute_value_loss(norm_s, actions)
219
+ loss_v.backward()
220
+ self.opt_v.step()
221
+
222
+ # Update Q
223
+ self.opt_q.zero_grad()
224
+ loss_q, q_mean = self.compute_q_loss(norm_s, actions, rewards, norm_s_next, dones)
225
+ loss_q.backward()
226
+ self.opt_q.step()
227
+
228
+ # Update Policy
229
+ self.opt_pi.zero_grad()
230
+ loss_pi, adv_mean, exp_adv_mean = self.compute_policy_loss(norm_s, actions)
231
+ loss_pi.backward()
232
+ self.opt_pi.step()
233
+
234
+ # Soft update target Q
235
+ soft_update(self.q1_target, self.q1, self.target_update_rate)
236
+ soft_update(self.q2_target, self.q2, self.target_update_rate)
237
+
238
+ return {
239
+ 'v_loss': loss_v.item(),
240
+ 'q_loss': loss_q.item(),
241
+ 'policy_loss': loss_pi.item(),
242
+ 'v_mean': v_mean,
243
+ 'q_mean': q_mean,
244
+ 'advantage_mean': adv_mean,
245
+ 'exp_advantage_mean': exp_adv_mean,
246
+ }
247
+
248
+ def save(self, path):
249
+ os.makedirs(os.path.dirname(path), exist_ok=True)
250
+ torch.save({
251
+ 'q1': self.q1.state_dict(),
252
+ 'q2': self.q2.state_dict(),
253
+ 'q1_target': self.q1_target.state_dict(),
254
+ 'q2_target': self.q2_target.state_dict(),
255
+ 'v': self.v.state_dict(),
256
+ 'policy': self.policy.state_dict(),
257
+ }, path)
258
+
259
+ def load(self, path):
260
+ ckpt = torch.load(path, map_location=self.device, weights_only=True)
261
+ self.q1.load_state_dict(ckpt['q1'])
262
+ self.q2.load_state_dict(ckpt['q2'])
263
+ self.q1_target.load_state_dict(ckpt['q1_target'])
264
+ self.q2_target.load_state_dict(ckpt['q2_target'])
265
+ self.v.load_state_dict(ckpt['v'])
266
+ self.policy.load_state_dict(ckpt['policy'])
267
+
268
+
269
+ # ============================================================
270
+ # Main training loop
271
+ # ============================================================
272
+
273
+ def train(tau_expectile=0.7, beta=3.0, lr=3e-4, batch_size=256,
274
+ discount=0.99, target_update_rate=0.005, num_iterations=100000,
275
+ seed=42, log_freq=1000, save_freq=10000):
276
+
277
+ # Seed
278
+ torch.manual_seed(seed)
279
+ np.random.seed(seed)
280
+
281
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
282
+ print(f"Device: {device}")
283
+
284
+ # Dataset
285
+ data_path = '/code/zxx240000/training/offline_rl/data/offline_dataset.npz'
286
+ dataset = OfflineDataset(data_path, device)
287
+
288
+ # Output dirs
289
+ tau_str = str(tau_expectile).replace('.', '')
290
+ result_dir = f'/code/zxx240000/training/offline_rl/results/iql_tau{tau_str}'
291
+ ckpt_dir = os.path.join(result_dir, 'checkpoints')
292
+ os.makedirs(ckpt_dir, exist_ok=True)
293
+
294
+ # CSV log
295
+ csv_path = os.path.join(result_dir, 'training_log.csv')
296
+ csv_fields = ['step', 'v_loss', 'q_loss', 'policy_loss',
297
+ 'v_mean', 'q_mean', 'advantage_mean', 'exp_advantage_mean']
298
+ csv_file = open(csv_path, 'w', newline='')
299
+ csv_writer = csv.DictWriter(csv_file, fieldnames=csv_fields)
300
+ csv_writer.writeheader()
301
+
302
+ # Trainer
303
+ trainer = IQLTrainer(dataset, device, tau_expectile=tau_expectile,
304
+ beta=beta, lr=lr, discount=discount,
305
+ target_update_rate=target_update_rate)
306
+
307
+ print(f"\nStarting IQL training: tau={tau_expectile}, beta={beta}, "
308
+ f"lr={lr}, batch={batch_size}, iters={num_iterations}")
309
+ print(f"Results dir: {result_dir}\n")
310
+
311
+ for step in range(1, num_iterations + 1):
312
+ metrics = trainer.train_step(batch_size)
313
+
314
+ if step % log_freq == 0:
315
+ row = {'step': step, **metrics}
316
+ csv_writer.writerow(row)
317
+ csv_file.flush()
318
+ print(f"[{step:>6d}] v_loss={metrics['v_loss']:.4f} "
319
+ f"q_loss={metrics['q_loss']:.4f} "
320
+ f"pi_loss={metrics['policy_loss']:.4f} "
321
+ f"V={metrics['v_mean']:.4f} Q={metrics['q_mean']:.4f} "
322
+ f"adv={metrics['advantage_mean']:.4f} "
323
+ f"exp_adv={metrics['exp_advantage_mean']:.2f}")
324
+
325
+ if step % save_freq == 0:
326
+ ckpt_path = os.path.join(ckpt_dir, f'iql_step{step}.pt')
327
+ trainer.save(ckpt_path)
328
+ print(f" -> Saved checkpoint: {ckpt_path}")
329
+
330
+ csv_file.close()
331
+
332
+ # Final save
333
+ final_path = os.path.join(ckpt_dir, 'iql_final.pt')
334
+ trainer.save(final_path)
335
+ print(f"\nTraining complete. Final checkpoint: {final_path}")
336
+
337
+ # Validation: sample actions and check limits
338
+ print("\n--- Validation ---")
339
+ norm_s_sample = dataset.normalize_state(dataset.states[:100])
340
+ with torch.no_grad():
341
+ sampled_actions, _ = trainer.policy.sample(norm_s_sample)
342
+ q1_vals = trainer.q1(norm_s_sample, dataset.actions[:100])
343
+ v_vals = trainer.v(norm_s_sample)
344
+
345
+ print(f"Sampled action min: {sampled_actions.min(dim=0).values.cpu().numpy()}")
346
+ print(f"Sampled action max: {sampled_actions.max(dim=0).values.cpu().numpy()}")
347
+ print(f"Joint lower limits: {JOINT_LOWER.cpu().numpy()}")
348
+ print(f"Joint upper limits: {JOINT_UPPER.cpu().numpy()}")
349
+ within = (sampled_actions >= JOINT_LOWER) & (sampled_actions <= JOINT_UPPER)
350
+ print(f"Actions within limits: {within.all().item()}")
351
+ print(f"Q-values range: [{q1_vals.min().item():.4f}, {q1_vals.max().item():.4f}]")
352
+ print(f"V-values range: [{v_vals.min().item():.4f}, {v_vals.max().item():.4f}]")
353
+ print(f"V < Q (expected): V_mean={v_vals.mean().item():.4f}, Q_mean={q1_vals.mean().item():.4f}")
354
+
355
+ # Verify checkpoint loads
356
+ test_trainer = IQLTrainer(dataset, device)
357
+ test_trainer.load(final_path)
358
+ with torch.no_grad():
359
+ test_v = test_trainer.v(norm_s_sample[:1])
360
+ print(f"Checkpoint reload test: V(s0) = {test_v.item():.4f} (should match)")
361
+
362
+ return trainer
363
+
364
+
365
+ if __name__ == '__main__':
366
+ parser = argparse.ArgumentParser()
367
+ parser.add_argument('--tau', type=float, default=0.7)
368
+ parser.add_argument('--beta', type=float, default=3.0)
369
+ parser.add_argument('--lr', type=float, default=3e-4)
370
+ parser.add_argument('--batch_size', type=int, default=256)
371
+ parser.add_argument('--discount', type=float, default=0.99)
372
+ parser.add_argument('--target_update_rate', type=float, default=0.005)
373
+ parser.add_argument('--num_iterations', type=int, default=100000)
374
+ parser.add_argument('--seed', type=int, default=42)
375
+ parser.add_argument('--log_freq', type=int, default=1000)
376
+ parser.add_argument('--save_freq', type=int, default=10000)
377
+ args = parser.parse_args()
378
+
379
+ train(
380
+ tau_expectile=args.tau,
381
+ beta=args.beta,
382
+ lr=args.lr,
383
+ batch_size=args.batch_size,
384
+ discount=args.discount,
385
+ target_update_rate=args.target_update_rate,
386
+ num_iterations=args.num_iterations,
387
+ seed=args.seed,
388
+ log_freq=args.log_freq,
389
+ save_freq=args.save_freq,
390
+ )
training_log.csv ADDED
@@ -0,0 +1,101 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ step,v_loss,q_loss,policy_loss,v_mean,q_mean,advantage_mean,exp_advantage_mean
2
+ 1000,0.0001488937414251268,0.017123185098171234,-17.457326889038086,0.04721087962388992,0.06451725959777832,-0.00599301652982831,0.9837942123413086
3
+ 2000,0.0005286661325953901,0.01033392921090126,-18.9333438873291,0.09802320599555969,0.11099211871623993,-0.006110389716923237,0.9875072240829468
4
+ 3000,0.0010939082130789757,0.006473910994827747,-19.31658172607422,0.1788613498210907,0.15808287262916565,-0.01295169722288847,0.9724383354187012
5
+ 4000,0.0012577481102198362,0.0070578427985310555,-19.839832305908203,0.21844390034675598,0.21006625890731812,-0.01451408863067627,0.9698702692985535
6
+ 5000,0.0006946508074179292,0.006777857430279255,-21.192907333374023,0.30034518241882324,0.2972208559513092,-0.013719955459237099,0.9666890501976013
7
+ 6000,0.0008185493643395603,0.00454492773860693,-21.101402282714844,0.3440312147140503,0.3355218172073364,-0.014879625290632248,0.9639429450035095
8
+ 7000,0.0011498236563056707,0.0043181260116398335,-22.2924747467041,0.3763347268104553,0.37749457359313965,-0.0027895004022866488,1.003147840499878
9
+ 8000,0.0005304793594405055,0.006958195939660072,-21.830453872680664,0.3898789882659912,0.3875240087509155,-0.009434308856725693,0.9775826930999756
10
+ 9000,0.0005791351431980729,0.009224031120538712,-22.123802185058594,0.4506993889808655,0.4588368535041809,-0.008604421280324459,0.979835033416748
11
+ 10000,0.000552843848709017,0.007071847096085548,-23.160566329956055,0.4691406488418579,0.46988046169281006,-0.007200233638286591,0.9843333959579468
12
+ 11000,0.0006002543959766626,0.005418932065367699,-21.51403045654297,0.5051866769790649,0.5040819644927979,-0.009896766394376755,0.9763743877410889
13
+ 12000,0.00021477842528838664,0.006183978170156479,-23.75772476196289,0.5187397599220276,0.5109261274337769,-0.005618714727461338,0.9855198860168457
14
+ 13000,0.00026278692530468106,0.004638044163584709,-23.86483383178711,0.5331025123596191,0.5372599363327026,-0.005588066764175892,0.9861727952957153
15
+ 14000,0.00040789341437630355,0.004485384561121464,-23.18863868713379,0.542675256729126,0.5440471768379211,-0.005035771988332272,0.9896504282951355
16
+ 15000,0.0004344815097283572,0.005826424807310104,-24.25287628173828,0.5674669742584229,0.555858850479126,-0.003824046580120921,0.9928783178329468
17
+ 16000,0.00030007955501787364,0.008056105114519596,-23.966318130493164,0.5731939673423767,0.5757041573524475,-0.005145515780895948,0.9877653121948242
18
+ 17000,0.0002703202480915934,0.005122414790093899,-24.526195526123047,0.5817747116088867,0.5819088816642761,-0.0036303987726569176,0.9921929240226746
19
+ 18000,0.00039316178299486637,0.00689063873142004,-23.439794540405273,0.5910090804100037,0.5759234428405762,-0.010704418644309044,0.9726033210754395
20
+ 19000,0.00018615866429172456,0.00721375085413456,-25.051410675048828,0.5937175154685974,0.5937061309814453,0.0011550185736268759,1.0055806636810303
21
+ 20000,0.00025072688004001975,0.0015508097130805254,-25.218236923217773,0.5873339176177979,0.5861520767211914,0.00019771966617554426,1.0030486583709717
22
+ 21000,0.00022000273747835308,0.0030086319893598557,-23.37228012084961,0.5993434190750122,0.6004638671875,-0.002670956775546074,0.9945750832557678
23
+ 22000,0.00017362178186886013,0.00239961757324636,-24.678878784179688,0.6005529165267944,0.596611499786377,-0.0016725929453969002,0.9968942403793335
24
+ 23000,0.00013084686361253262,0.006569467484951019,-24.938079833984375,0.6206008195877075,0.6142706871032715,-0.0033857692033052444,0.9911432266235352
25
+ 24000,0.00018223526421934366,0.002369912341237068,-24.460426330566406,0.6045263409614563,0.5985832810401917,-0.005711064673960209,0.985088586807251
26
+ 25000,0.00010868997196666896,0.004901471547782421,-25.510717391967773,0.6051712036132812,0.6084377765655518,-3.5031349398195744e-05,1.001131534576416
27
+ 26000,0.00019300405983813107,0.00422974769026041,-26.030649185180664,0.6115922927856445,0.6020850539207458,-0.006489139050245285,0.9829984903335571
28
+ 27000,0.00020204705651849508,0.006951279006898403,-25.439002990722656,0.6009725332260132,0.5989433526992798,0.0013640759279951453,1.0060769319534302
29
+ 28000,0.00021440471755340695,0.0036399043165147305,-25.49695587158203,0.6086745858192444,0.6052184104919434,0.0013740079011768103,1.0065057277679443
30
+ 29000,8.494644134771079e-05,0.004319016356021166,-25.657848358154297,0.6075325608253479,0.6071163415908813,-0.0016160481609404087,0.9961751103401184
31
+ 30000,7.651586929569021e-05,0.004498209338635206,-25.5654354095459,0.6022244095802307,0.6025192737579346,-0.0028532766737043858,0.9923875331878662
32
+ 31000,0.00014085255679674447,0.0028083024080842733,-25.80255889892578,0.600712776184082,0.6001958250999451,-0.0006916482234373689,0.9995920062065125
33
+ 32000,0.00012024045281577855,0.006574017461389303,-24.54662322998047,0.5963567495346069,0.5882075428962708,-0.005151441786438227,0.9860668182373047
34
+ 33000,0.00015378290845546871,0.011917476542294025,-25.39267921447754,0.5959603190422058,0.5814091563224792,-0.004795705899596214,0.9874920845031738
35
+ 34000,0.00019843160407617688,0.0028308317996561527,-25.03889274597168,0.6032512187957764,0.6092047095298767,0.0007185174617916346,1.004421591758728
36
+ 35000,0.00013722555013373494,0.006110752001404762,-25.946321487426758,0.5967638492584229,0.5925816297531128,-0.004547523334622383,0.9881558418273926
37
+ 36000,0.00011616927804425359,0.006658279802650213,-25.817729949951172,0.5974419116973877,0.589691162109375,-0.008166152983903885,0.9768060445785522
38
+ 37000,0.00015622015052940696,0.00927338469773531,-26.523395538330078,0.5916243195533752,0.5916246175765991,-0.000850978191010654,0.9992008209228516
39
+ 38000,0.00016840218449942768,0.0005011010216549039,-25.401657104492188,0.5957861542701721,0.5887603759765625,-0.00331634609028697,0.9920378923416138
40
+ 39000,0.00010988415306201205,0.007036756724119186,-26.540027618408203,0.587198793888092,0.5874839425086975,-0.003404047805815935,0.9908913969993591
41
+ 40000,0.00019101548241451383,0.010661423206329346,-25.81806182861328,0.5920298099517822,0.5750237703323364,-0.005834830924868584,0.9848833084106445
42
+ 41000,0.00013355040573514998,0.002639091107994318,-25.92420196533203,0.5998861789703369,0.592593789100647,-0.005590290296822786,0.984661340713501
43
+ 42000,0.00013216228398960084,0.0030659311451017857,-26.353906631469727,0.5818030834197998,0.5826287865638733,-0.0015514675760641694,0.9969161748886108
44
+ 43000,0.00012512158718891442,0.0013485478702932596,-25.845565795898438,0.5800286531448364,0.5821515321731567,-0.005123673472553492,0.9859508275985718
45
+ 44000,0.00012281336239539087,0.009433717466890812,-26.624422073364258,0.5814752578735352,0.5757251977920532,-0.0027976958081126213,0.9930979609489441
46
+ 45000,0.00010371934331487864,0.009669829159975052,-27.00676727294922,0.5739758014678955,0.5833988189697266,0.0028173085302114487,1.0097479820251465
47
+ 46000,9.72836569417268e-05,0.005585405509918928,-26.44808006286621,0.5774369239807129,0.5724466443061829,-0.0031470314133912325,0.9917936325073242
48
+ 47000,9.05219349078834e-05,0.008765907026827335,-26.231090545654297,0.5737698078155518,0.5855615735054016,-5.775375757366419e-05,1.000643253326416
49
+ 48000,0.00014932529302313924,0.0039146654307842255,-25.868698120117188,0.5787293314933777,0.5797147750854492,-0.00457192724570632,0.9880678653717041
50
+ 49000,0.00015518236614298075,0.0016166161512956023,-25.75746726989746,0.5819497108459473,0.5806148052215576,-0.00766854640096426,0.9790775775909424
51
+ 50000,9.765935828909278e-05,0.005000343546271324,-26.063396453857422,0.5749824047088623,0.5687891244888306,-0.004873913247138262,0.9865258932113647
52
+ 51000,6.426209438359365e-05,0.0021728933788836002,-25.74820327758789,0.5731598138809204,0.5830702781677246,-0.0007245639571920037,0.9981404542922974
53
+ 52000,5.935573426540941e-05,0.003119529690593481,-26.813621520996094,0.5822153091430664,0.5772654414176941,-0.004683001898229122,0.9867243766784668
54
+ 53000,0.00013016824959777296,0.003621558891609311,-26.15018081665039,0.5665186047554016,0.5619409084320068,-0.003845440223813057,0.9900767207145691
55
+ 54000,5.861584941158071e-05,0.0019813221879303455,-26.52945327758789,0.5710108280181885,0.5725727081298828,-0.00014177733100950718,1.000206708908081
56
+ 55000,0.00012967182556167245,0.0027217124588787556,-26.197147369384766,0.5730605125427246,0.5763099193572998,-0.00575302867218852,0.9843593835830688
57
+ 56000,5.904484714847058e-05,0.002086964203044772,-26.71909523010254,0.5752797722816467,0.5784851312637329,-0.0036513577215373516,0.9895596504211426
58
+ 57000,0.0001286962942685932,0.0025332202203571796,-27.35590171813965,0.5737621784210205,0.5751775503158569,0.0008270799880847335,1.004018783569336
59
+ 58000,7.815567369107157e-05,0.006133121903985739,-26.280715942382812,0.5805566906929016,0.5763545632362366,-0.004147685598582029,0.9885240793228149
60
+ 59000,7.395993452519178e-05,0.008265103213489056,-26.150333404541016,0.5733118057250977,0.5674525499343872,-0.0022185463458299637,0.9940261244773865
61
+ 60000,8.750478446017951e-05,0.001516152173280716,-26.9849853515625,0.5730513334274292,0.5679193735122681,-0.004655994474887848,0.9871608018875122
62
+ 61000,0.00011709539830917493,0.0008902038680389524,-25.748294830322266,0.5814056992530823,0.5742424726486206,-0.006016991101205349,0.9833992123603821
63
+ 62000,0.0001369673991575837,0.007829621434211731,-26.747182846069336,0.5807332992553711,0.5711119174957275,-0.00680540269240737,0.9812262058258057
64
+ 63000,0.00011835618352051824,0.01048852689564228,-26.834095001220703,0.5697199702262878,0.5688411593437195,-0.002711495151743293,0.9930967688560486
65
+ 64000,0.00010618756641633809,0.0011887461878359318,-25.574569702148438,0.5720807313919067,0.5710684061050415,-0.0034998124465346336,0.99074387550354
66
+ 65000,0.00010482417565071955,0.005848865956068039,-26.042343139648438,0.5716762542724609,0.5632674098014832,-0.004522524308413267,0.9877628087997437
67
+ 66000,8.895703649614006e-05,0.006089819595217705,-27.129566192626953,0.560055136680603,0.5623401403427124,-0.0003869016654789448,0.9999047517776489
68
+ 67000,0.00011330326378811151,0.006693074479699135,-23.250038146972656,0.5589816570281982,0.5644457340240479,-0.0018776217475533485,0.9955991506576538
69
+ 68000,9.079788287635893e-05,0.0031706918962299824,-26.634340286254883,0.5578062534332275,0.5601974129676819,-0.002269890159368515,0.9943193197250366
70
+ 69000,5.73778452235274e-05,0.006508987862616777,-26.516582489013672,0.5657087564468384,0.5613890886306763,-0.007151088677346706,0.9794530272483826
71
+ 70000,8.407401764998212e-05,0.0019406850915402174,-27.36440658569336,0.5493730902671814,0.5476880669593811,-0.003952967002987862,0.9890477657318115
72
+ 71000,0.00016143931134138256,0.006198926828801632,-26.58662223815918,0.5561314821243286,0.5490660667419434,-0.007141285575926304,0.9803811311721802
73
+ 72000,8.3572042058222e-05,0.0032507595606148243,-26.001232147216797,0.555461049079895,0.549469530582428,-0.005260828882455826,0.9850512146949768
74
+ 73000,0.00010661175474524498,0.0007440145127475262,-26.81865882873535,0.5565646886825562,0.5539497137069702,-0.0019408157095313072,0.9954901933670044
75
+ 74000,0.00024070873041637242,0.0052931467071175575,-27.776180267333984,0.5686681866645813,0.5621130466461182,-0.007694175001233816,0.9798183441162109
76
+ 75000,0.00015474247629754245,0.005048967897891998,-27.402301788330078,0.5500936508178711,0.5413309335708618,-0.002668903209269047,0.9936968088150024
77
+ 76000,7.424000796163455e-05,0.0038600116968154907,-25.61629295349121,0.5703386068344116,0.5637786388397217,-0.002474918495863676,0.9934521913528442
78
+ 77000,8.965870074462146e-05,0.006210260558873415,-27.60029411315918,0.5523503422737122,0.5520658493041992,-0.004888276569545269,0.9864134788513184
79
+ 78000,0.00010770530207082629,0.004852974787354469,-27.586360931396484,0.5570819973945618,0.5587167143821716,-0.0017596060642972589,0.9959943294525146
80
+ 79000,0.00010430492693558335,0.0048827421851456165,-27.83783721923828,0.5605653524398804,0.5653701424598694,-0.0019765314646065235,0.9948995113372803
81
+ 80000,0.00010558339999988675,0.00285049295052886,-26.190105438232422,0.5531935691833496,0.5483263731002808,-0.002991686575114727,0.9923228025436401
82
+ 81000,8.737208554521203e-05,0.0008890702156350017,-26.73209571838379,0.5547089576721191,0.5558770895004272,-0.0003758469829335809,0.9996782541275024
83
+ 82000,8.134762174449861e-05,0.004000999964773655,-26.78078842163086,0.5528160929679871,0.5524387955665588,-0.0018687330884858966,0.9953027963638306
84
+ 83000,0.00012431899085640907,0.005905299913138151,-27.33001136779785,0.5564279556274414,0.556323766708374,-0.00594191811978817,0.9836968183517456
85
+ 84000,0.00020311625848989934,0.005271616391837597,-27.606876373291016,0.5467767715454102,0.545536994934082,-0.0037131914868950844,0.9911004304885864
86
+ 85000,6.36972690699622e-05,0.001846288563683629,-26.49146270751953,0.554013729095459,0.5509346127510071,-0.003998124040663242,0.9887819886207581
87
+ 86000,7.593004556838423e-05,0.008715381845831871,-27.499427795410156,0.5444061756134033,0.5431278944015503,-0.002713325899094343,0.9928195476531982
88
+ 87000,0.00011334389273542911,0.0015394220827147365,-26.958168029785156,0.5473346710205078,0.5436162948608398,-0.0051136803813278675,0.9861065149307251
89
+ 88000,8.364733366761357e-05,0.0004284453170839697,-26.754413604736328,0.5363893508911133,0.533919632434845,-0.003050466999411583,0.9918617010116577
90
+ 89000,8.446662104688585e-05,0.002174799796193838,-26.83147430419922,0.5341291427612305,0.5379350185394287,-0.003239447483792901,0.9912080764770508
91
+ 90000,7.296711555682123e-05,0.0037364636082202196,-27.3741397857666,0.5361378192901611,0.5344648957252502,-0.001979462569579482,0.9947989583015442
92
+ 91000,0.0001922717347042635,0.002574441023170948,-26.53437614440918,0.5284481048583984,0.5287431478500366,-0.004611331969499588,0.9885634779930115
93
+ 92000,8.725952648092061e-05,0.0029895692132413387,-27.269290924072266,0.5517081618309021,0.5533859729766846,-0.00309176673181355,0.9917852878570557
94
+ 93000,7.919182098703459e-05,0.0047914437018334866,-27.147789001464844,0.5381056070327759,0.5327804088592529,-0.004360011778771877,0.9878666996955872
95
+ 94000,0.00018320753588341177,0.002686972264200449,-26.622901916503906,0.5366100072860718,0.5350701212882996,-0.004095324315130711,0.9899418354034424
96
+ 95000,0.00014682550681754947,0.0029815668240189552,-27.30376625061035,0.5415458679199219,0.535621702671051,-0.0025409297086298466,0.9939871430397034
97
+ 96000,0.00012204082304378971,0.010050460696220398,-28.051822662353516,0.5348919034004211,0.537726879119873,-0.0016675336519256234,0.9962784051895142
98
+ 97000,0.00019438120943959802,0.0029506103601306677,-26.26782989501953,0.5408000946044922,0.5382518172264099,-0.004446289502084255,0.9889114499092102
99
+ 98000,0.00012340216198936105,0.0007030363194644451,-27.668563842773438,0.5397628545761108,0.5373809337615967,-0.0035567376762628555,0.9906792640686035
100
+ 99000,0.00011772227298934013,0.006932809017598629,-27.143587112426758,0.5412322282791138,0.5380535125732422,-0.000751726794987917,0.9989335536956787
101
+ 100000,0.000119170275866054,0.004066753666847944,-27.553874969482422,0.5407698154449463,0.5363407135009766,-0.0040032402612268925,0.9894227981567383