tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it./app/objectives.py
1 from __future__ import annotations
2
3 import torch
4
5
6 def masked_mean(values, mask, axis=None):
7 raise NotImplementedError
8
9
10 def masked_sum(values, mask, axis=None):
11 raise NotImplementedError
12
13
14 def logsumexp(x, axis):
15 raise NotImplementedError
16
17
18 def log_softmax(x, axis):
19 raise NotImplementedError
20
21
22 def selective_logprobs(logits, labels, mask):
23 raise NotImplementedError
24
25
26 def token_logprobs(logits, labels):
27 raise NotImplementedError
28
29
30 def sequence_logprob(logits, labels, mask, length_normalize):
31 raise NotImplementedError
32
33
34 def entropy(logits, mask):
35 raise NotImplementedError
36
37
38 def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
39 raise NotImplementedError
40
41
42 def ipo_loss(pc, pr, rc, rr, beta):
43 raise NotImplementedError
44
45
46 def grpo_advantages(rewards, group_size, scale_by_std):
47 raise NotImplementedError
48
49
50 def gae(rewards, values, next_value, gamma, lam):
51 raise NotImplementedError
52
53
54 def kl_penalty(logp, ref_logp, estimator):
55 raise NotImplementedError
56
57
58 def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
59 raise NotImplementedError
60
61
62 def value_loss(values, old_values, returns, clip):
63 raise NotImplementedError
64
65
66 def whiten(values, mask, shift_mean):
67 raise NotImplementedError
68
69
70 def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
71 chosen_labels, rejected_labels, chosen_mask, rejected_mask,
72 beta, label_smoothing):
73 raise NotImplementedError
74
75
76 def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
77 rewards, group_size, beta, clip_low, clip_high, scale_by_std,
78 kl_estimator):
79 raise NotImplementedError
80
81
82 def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
83 gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
84 raise NotImplementedError
85
86
87 def rloo_advantages(rewards, group_size):
88 raise NotImplementedError
89
90
91 def reverse_kl(logp, ref_logp):
92 raise NotImplementedError
93
94
95 def importance_ratio(logp, old_logp, clip):
96 raise NotImplementedError
97
98
99 def discounted_returns(rewards, gamma):
100 raise NotImplementedError
101
102
103 def normalize(x, eps):
104 raise NotImplementedError
105
106
107 def top_p_mask(probs, p):
108 raise NotImplementedError
109
110
111 def smoothed_nll(logits, labels, smoothing):
112 raise NotImplementedError
113
114
115 def bradley_terry_logit(chosen_reward, rejected_reward, beta):
116 raise NotImplementedError
117
118
119 def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
120 raise NotImplementedError
121
122
123 def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
124 raise NotImplementedError
125
126
127 def cross_entropy(logits, labels, ignore_index):
128 raise NotImplementedError
129
130
131 def top_k_mask(logits, k):
132 raise NotImplementedError
133
134
135 def group_mean_baseline(rewards, group_size):
136 raise NotImplementedError
137
138
139 def lambda_returns(rewards, values, next_value, gamma, lam):
140 raise NotImplementedError
141
142
143 def symmetric_kl(logp, ref_logp):
144 raise NotImplementedError
145
146
147 def huber_value_loss(values, returns, delta):
148 raise NotImplementedError
149
150
151 def normalized_entropy(logits, mask):
152 raise NotImplementedError
153
154
155 def clip_fraction(logp, old_logp, clip):
156 raise NotImplementedError
157
158
159 def masked_whiten(values, mask, shift_mean):
160 raise NotImplementedError
161
162
163 def logprob_at_temperature(logits, labels, temperature):
164 raise NotImplementedError
165
166
167 def advantage_mean_std(advantages, mask):
168 raise NotImplementedError
169
170
171 def argmax_tokens(logits):
172 raise NotImplementedError
173
174
175 def mode_label(labels):
176 raise NotImplementedError
177
178
179 def median_reward(rewards):
180 raise NotImplementedError
181
182
183 def masked_max(values, mask):
184 raise NotImplementedError
185
186
187 def quantile_lower(x, q):
188 raise NotImplementedError
189
190
191 def pad_mask_from_lengths(lengths, max_len):
192 raise NotImplementedError
193
194
195 def first_nonzero_index(mask):
196 raise NotImplementedError
197
198
199 def cumulative_max(x):
200 raise NotImplementedError
201
202
203 def bucketize_reward(rewards, boundaries):
204 raise NotImplementedError
205
1 from __future__ import annotations
2
3 import torch
4
5 import reforacle
6
7 g = torch.Generator().manual_seed(0)
8 values = torch.randn(8, generator=g)
9 mask = (torch.rand(8, generator=g) > 0.5).float()
10
11 print("inputs:")
12 print(" values =", values.tolist())
13 print(" mask =", mask.tolist())
14 print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
15 print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
16 print()
17 print("Probe any reference function this way to discover its exact numerical")
18 print("convention, then implement it in /app/objectives.py to match.")
19
Remaining time (hours:minutes):
1:59
[stdout]
Remaining time (hours:minutes):
1:59/app/probe.py
1 from __future__ import annotations
2
3 import torch
4
5
6 def masked_mean(values, mask, axis=None):
7 raise NotImplementedError
8
9
10 def masked_sum(values, mask, axis=None):
11 raise NotImplementedError
12
13
14 def logsumexp(x, axis):
15 raise NotImplementedError
16
17
18 def log_softmax(x, axis):
19 raise NotImplementedError
20
21
22 def selective_logprobs(logits, labels, mask):
23 raise NotImplementedError
24
25
26 def token_logprobs(logits, labels):
27 raise NotImplementedError
28
29
30 def sequence_logprob(logits, labels, mask, length_normalize):
31 raise NotImplementedError
32
33
34 def entropy(logits, mask):
35 raise NotImplementedError
36
37
38 def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
39 raise NotImplementedError
40
41
42 def ipo_loss(pc, pr, rc, rr, beta):
43 raise NotImplementedError
44
45
46 def grpo_advantages(rewards, group_size, scale_by_std):
47 raise NotImplementedError
48
49
50 def gae(rewards, values, next_value, gamma, lam):
51 raise NotImplementedError
52
53
54 def kl_penalty(logp, ref_logp, estimator):
55 raise NotImplementedError
56
57
58 def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
59 raise NotImplementedError
60
61
62 def value_loss(values, old_values, returns, clip):
63 raise NotImplementedError
64
65
66 def whiten(values, mask, shift_mean):
67 raise NotImplementedError
68
69
70 def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
71 chosen_labels, rejected_labels, chosen_mask, rejected_mask,
72 beta, label_smoothing):
73 raise NotImplementedError
74
75
76 def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
77 rewards, group_size, beta, clip_low, clip_high, scale_by_std,
78 kl_estimator):
79 raise NotImplementedError
80
81
82 def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
83 gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
84 raise NotImplementedError
85
86
87 def rloo_advantages(rewards, group_size):
88 raise NotImplementedError
89
90
91 def reverse_kl(logp, ref_logp):
92 raise NotImplementedError
93
94
95 def importance_ratio(logp, old_logp, clip):
96 raise NotImplementedError
97
98
99 def discounted_returns(rewards, gamma):
100 raise NotImplementedError
101
102
103 def normalize(x, eps):
104 raise NotImplementedError
105
106
107 def top_p_mask(probs, p):
108 raise NotImplementedError
109
110
111 def smoothed_nll(logits, labels, smoothing):
112 raise NotImplementedError
113
114
115 def bradley_terry_logit(chosen_reward, rejected_reward, beta):
116 raise NotImplementedError
117
118
119 def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
120 raise NotImplementedError
121
122
123 def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
124 raise NotImplementedError
125
126
127 def cross_entropy(logits, labels, ignore_index):
128 raise NotImplementedError
129
130
131 def top_k_mask(logits, k):
132 raise NotImplementedError
133
134
135 def group_mean_baseline(rewards, group_size):
136 raise NotImplementedError
137
138
139 def lambda_returns(rewards, values, next_value, gamma, lam):
140 raise NotImplementedError
141
142
143 def symmetric_kl(logp, ref_logp):
144 raise NotImplementedError
145
146
147 def huber_value_loss(values, returns, delta):
148 raise NotImplementedError
149
150
151 def normalized_entropy(logits, mask):
152 raise NotImplementedError
153
154
155 def clip_fraction(logp, old_logp, clip):
156 raise NotImplementedError
157
158
159 def masked_whiten(values, mask, shift_mean):
160 raise NotImplementedError
161
162
163 def logprob_at_temperature(logits, labels, temperature):
164 raise NotImplementedError
165
166
167 def advantage_mean_std(advantages, mask):
168 raise NotImplementedError
169
170
171 def argmax_tokens(logits):
172 raise NotImplementedError
173
174
175 def mode_label(labels):
176 raise NotImplementedError
177
178
179 def median_reward(rewards):
180 raise NotImplementedError
181
182
183 def masked_max(values, mask):
184 raise NotImplementedError
185
186
187 def quantile_lower(x, q):
188 raise NotImplementedError
189
190
191 def pad_mask_from_lengths(lengths, max_len):
192 raise NotImplementedError
193
194
195 def first_nonzero_index(mask):
196 raise NotImplementedError
197
198
199 def cumulative_max(x):
200 raise NotImplementedError
201
202
203 def bucketize_reward(rewards, boundaries):
204 raise NotImplementedError
205
1 from __future__ import annotations
2
3 import torch
4
5 import reforacle
6
7 g = torch.Generator().manual_seed(0)
8 values = torch.randn(8, generator=g)
9 mask = (torch.rand(8, generator=g) > 0.5).float()
10
11 print("inputs:")
12 print(" values =", values.tolist())
13 print(" mask =", mask.tolist())
14 print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
15 print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
16 print()
17 print("Probe any reference function this way to discover its exact numerical")
18 print("convention, then implement it in /app/objectives.py to match.")
19
Remaining time (hours:minutes):
1:59
[stdout]
Remaining time (hours:minutes):
1:59Check remaining budget
bash /app/timer.sh
1 from __future__ import annotations
2
3 import torch
4
5
6 def masked_mean(values, mask, axis=None):
7 raise NotImplementedError
8
9
10 def masked_sum(values, mask, axis=None):
11 raise NotImplementedError
12
13
14 def logsumexp(x, axis):
15 raise NotImplementedError
16
17
18 def log_softmax(x, axis):
19 raise NotImplementedError
20
21
22 def selective_logprobs(logits, labels, mask):
23 raise NotImplementedError
24
25
26 def token_logprobs(logits, labels):
27 raise NotImplementedError
28
29
30 def sequence_logprob(logits, labels, mask, length_normalize):
31 raise NotImplementedError
32
33
34 def entropy(logits, mask):
35 raise NotImplementedError
36
37
38 def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
39 raise NotImplementedError
40
41
42 def ipo_loss(pc, pr, rc, rr, beta):
43 raise NotImplementedError
44
45
46 def grpo_advantages(rewards, group_size, scale_by_std):
47 raise NotImplementedError
48
49
50 def gae(rewards, values, next_value, gamma, lam):
51 raise NotImplementedError
52
53
54 def kl_penalty(logp, ref_logp, estimator):
55 raise NotImplementedError
56
57
58 def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
59 raise NotImplementedError
60
61
62 def value_loss(values, old_values, returns, clip):
63 raise NotImplementedError
64
65
66 def whiten(values, mask, shift_mean):
67 raise NotImplementedError
68
69
70 def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
71 chosen_labels, rejected_labels, chosen_mask, rejected_mask,
72 beta, label_smoothing):
73 raise NotImplementedError
74
75
76 def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
77 rewards, group_size, beta, clip_low, clip_high, scale_by_std,
78 kl_estimator):
79 raise NotImplementedError
80
81
82 def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
83 gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
84 raise NotImplementedError
85
86
87 def rloo_advantages(rewards, group_size):
88 raise NotImplementedError
89
90
91 def reverse_kl(logp, ref_logp):
92 raise NotImplementedError
93
94
95 def importance_ratio(logp, old_logp, clip):
96 raise NotImplementedError
97
98
99 def discounted_returns(rewards, gamma):
100 raise NotImplementedError
101
102
103 def normalize(x, eps):
104 raise NotImplementedError
105
106
107 def top_p_mask(probs, p):
108 raise NotImplementedError
109
110
111 def smoothed_nll(logits, labels, smoothing):
112 raise NotImplementedError
113
114
115 def bradley_terry_logit(chosen_reward, rejected_reward, beta):
116 raise NotImplementedError
117
118
119 def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
120 raise NotImplementedError
121
122
123 def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
124 raise NotImplementedError
125
126
127 def cross_entropy(logits, labels, ignore_index):
128 raise NotImplementedError
129
130
131 def top_k_mask(logits, k):
132 raise NotImplementedError
133
134
135 def group_mean_baseline(rewards, group_size):
136 raise NotImplementedError
137
138
139 def lambda_returns(rewards, values, next_value, gamma, lam):
140 raise NotImplementedError
141
142
143 def symmetric_kl(logp, ref_logp):
144 raise NotImplementedError
145
146
147 def huber_value_loss(values, returns, delta):
148 raise NotImplementedError
149
150
151 def normalized_entropy(logits, mask):
152 raise NotImplementedError
153
154
155 def clip_fraction(logp, old_logp, clip):
156 raise NotImplementedError
157
158
159 def masked_whiten(values, mask, shift_mean):
160 raise NotImplementedError
161
162
163 def logprob_at_temperature(logits, labels, temperature):
164 raise NotImplementedError
165
166
167 def advantage_mean_std(advantages, mask):
168 raise NotImplementedError
169
170
171 def argmax_tokens(logits):
172 raise NotImplementedError
173
174
175 def mode_label(labels):
176 raise NotImplementedError
177
178
179 def median_reward(rewards):
180 raise NotImplementedError
181
182
183 def masked_max(values, mask):
184 raise NotImplementedError
185
186
187 def quantile_lower(x, q):
188 raise NotImplementedError
189
190
191 def pad_mask_from_lengths(lengths, max_len):
192 raise NotImplementedError
193
194
195 def first_nonzero_index(mask):
196 raise NotImplementedError
197
198
199 def cumulative_max(x):
200 raise NotImplementedError
201
202
203 def bucketize_reward(rewards, boundaries):
204 raise NotImplementedError
205
1 from __future__ import annotations
2
3 import torch
4
5 import reforacle
6
7 g = torch.Generator().manual_seed(0)
8 values = torch.randn(8, generator=g)
9 mask = (torch.rand(8, generator=g) > 0.5).float()
10
11 print("inputs:")
12 print(" values =", values.tolist())
13 print(" mask =", mask.tolist())
14 print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
15 print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
16 print()
17 print("Probe any reference function this way to discover its exact numerical")
18 print("convention, then implement it in /app/objectives.py to match.")
19
Remaining time (hours:minutes):
1:59
[stdout]
Remaining time (hours:minutes):
1:59List reforacle attributes
python3 -c "
import reforacle
fns = [x for x in dir(reforacle) if not x.startswith('_')]
print(len(fns))
print(fns)
"51 ['F', 'advantage_mean_std', 'argmax_tokens', 'bradley_terry_logit', 'bucketize_reward', 'clip_fraction', 'clipped_pg_loss', 'cross_entropy', 'cumulative_max', 'discounted_returns', 'dpo_loss', 'dpo_sequence_loss', 'entropy', 'first_nonzero_index', 'gae', 'group_mean_baseline', 'grpo_advantages', 'grpo_objective', 'huber_value_loss', 'importance_ratio', 'ipo_loss', 'kl_penalty', 'lambda_returns', 'log_softmax', 'logprob_at_temperature', 'logsumexp', 'masked_max', 'masked_mean', 'masked_sum', 'masked_whiten', 'median_reward', 'mode_label', 'normalize', 'normalized_entropy', 'pad_mask_from_lengths', 'ppo_objective', 'quantile_lower', 'reverse_kl', 'reverse_kl_objective', 'rloo_advantages', 'rloo_objective', 'selective_logprobs', 'sequence_logprob', 'smoothed_nll', 'symmetric_kl', 'token_logprobs', 'top_k_mask', 'top_p_mask', 'torch', 'value_loss', 'whiten'] [stdout] 51 ['F', 'advantage_mean_std', 'argmax_tokens', 'bradley_terry_logit', 'bucketize_reward', 'clip_fraction', 'clipped_pg_loss', 'cross_entropy', 'cumulative_max', 'discounted_returns', 'dpo_loss', 'dpo_sequence_loss', 'entropy', 'first_nonzero_index', 'gae', 'group_mean_baseline', 'grpo_advantages', 'grpo_objective', 'huber_value_loss', 'importance_ratio', 'ipo_loss', 'kl_penalty', 'lambda_returns', 'log_softmax', 'logprob_at_temperature', 'logsumexp', 'masked_max', 'masked_mean', 'masked_sum', 'masked_whiten', 'median_reward', 'mode_label', 'normalize', 'normalized_entropy', 'pad_mask_from_lengths', 'ppo_objective', 'quantile_lower', 'reverse_kl', 'reverse_kl_objective', 'rloo_advantages', 'rloo_objective', 'selective_logprobs', 'sequence_logprob', 'smoothed_nll', 'symmetric_kl', 'token_logprobs', 'top_k_mask', 'top_p_mask', 'torch', 'value_loss', 'whiten']
Probe masked_mean/sum
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(1)
def show(name,*a):
try: print(name,'->',a[0])
except: pass
# masked_mean / masked_sum
v=torch.randn(3,4,generator=g); m=(torch.rand(3,4,generator=g)>0.4).float()
print('masked_mean none:', R.masked_mean(v,m).item())
print('masked_mean ax0:', R.masked_mean(v,m,axis=0).tolist())
print('masked_mean ax1:', R.masked_mean(v,m,axis=1).tolist())
print('masked_sum none:', R.masked_sum(v,m).item())
print('masked_sum ax1:', R.masked_sum(v,m,axis=1).tolist())
# empty axis mean -> 0/0?
m2=m.clone(); m2[0]=0
print('masked_mean ax1 empty row:', R.masked_mean(v,m2,axis=1).tolist())
print('masked_mean none check manual:', (v*m).sum().item()/m.sum().item())
"masked_mean none: -0.1978435516357422 masked_mean ax0: [-1.0276086330413818, 0.05039694160223007, -0.7305456399917603, 0.5015006065368652] masked_mean ax1: [0.31663957238197327, -0.4357382357120514, -1.0276086330413818] masked_sum none: -1.3849048614501953 masked_sum ax1: [0.9499186873435974, -1.3072147369384766, -1.0276086330413818] masked_mean ax1 empty row: [0.0, -0.4357382357120514, -1.0276086330413818] masked_mean none check manual: -0.1978435516357422 [stdout] masked_mean none: -0.1978435516357422 masked_mean ax0: [-1.0276086330413818, 0.05039694160223007, -0.7305456399917603, 0.5015006065368652] masked_mean ax1: [0.31663957238197327, -0.4357382357120514, -1.0276086330413818] masked_sum none: -1.3849048614501953 masked_sum ax1: [0.9499186873435974, -1.3072147369384766, -1.0276086330413818] masked_mean ax1 empty row: [0.0, -0.4357382357120514, -1.0276086330413818] masked_mean none check manual: -0.1978435516357422
Probe logsumexp, log_softmax, normalize
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(2)
x=torch.randn(3,4,generator=g)
print('logsumexp ax1:', R.logsumexp(x,1).tolist())
print(' torch :', torch.logsumexp(x,1).tolist())
print('log_softmax ax1 row0:', R.log_softmax(x,1)[0].tolist())
print(' torch :', torch.log_softmax(x,1)[0].tolist())
print('normalize:', R.normalize(x, 1e-5)[0].tolist())
print(' manual:', ((x-x.mean())/(x.std()+1e-5))[0].tolist())
print(' manual2:', ((x-x.mean())/(x.std(unbiased=False)+1e-5))[0].tolist())
"logsumexp ax1: [1.1958097219467163, 1.9263124465942383, 1.928405523300171] torch : [1.1958097219467163, 1.9263124465942383, 1.928405523300171] log_softmax ax1 row0: [-0.8035129308700562, -1.4193737506866455, -1.5153100490570068, -2.4008469581604004] torch : [-0.8035128712654114, -1.4193737506866455, -1.5153100490570068, -2.4008467197418213] normalize: [0.3980754315853119, -0.3550233542919159, -0.4723379909992218, -1.5552071332931519] manual: [0.3980754315853119, -0.3550233542919159, -0.4723379909992218, -1.5552071332931519] manual2: [0.41577601432800293, -0.3708096146583557, -0.4933406710624695, -1.6243599653244019] [stdout] logsumexp ax1: [1.1958097219467163, 1.9263124465942383, 1.928405523300171] torch : [1.1958097219467163, 1.9263124465942383, 1.928405523300171] log_softmax ax1 row0: [-0.8035129308700562, -1.4193737506866455, -1.5153100490570068, -2.4008469581604004] torch : [-0.8035128712654114, -1.4193737506866455, -1.5153100490570068, -2.4008467197418213] normalize: [0.3980754315853119, -0.3550233542919159, -0.4723379909992218, -1.5552071332931519] manual: [0.3980754315853119, -0.3550233542919159, -0.4723379909992218, -1.5552071332931519] manual2: [0.41577601432800293, -0.3708096146583557, -0.4933406710624695, -1.6243599653244019]
Probe token/selective/sequence logprobs
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(3)
# token_logprobs / selective_logprobs / sequence_logprob
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
lp=torch.log_softmax(logits,-1)
tl=torch.gather(lp,-1,labels.unsqueeze(-1)).squeeze(-1)
print('token_logprobs:', R.token_logprobs(logits,labels).tolist())
print(' manual :', tl.tolist())
print('selective_logprobs:', R.selective_logprobs(logits,labels,mask).tolist())
print(' masked tl :', (tl*mask).tolist())
print('sequence_logprob norm=False:', R.sequence_logprob(logits,labels,mask,False).tolist())
print(' manual sum:', (tl*mask).sum(-1).tolist())
print('sequence_logprob norm=True:', R.sequence_logprob(logits,labels,mask,True).tolist())
print(' manual sum/masksum:', ((tl*mask).sum(-1)/mask.sum(-1)).tolist())
"token_logprobs: [[-1.4040248394012451, -2.092271089553833, -0.9929482340812683], [-1.7721467018127441, -1.8669283390045166, -1.6138648986816406]] manual : [[-1.4040248394012451, -2.092271089553833, -0.9929482340812683], [-1.7721467018127441, -1.8669283390045166, -1.6138648986816406]] selective_logprobs: [-4.489243984222412, -5.2529401779174805] masked tl : [[-1.4040248394012451, -2.092271089553833, -0.9929482340812683], [-1.7721467018127441, -1.8669283390045166, -1.6138648986816406]] sequence_logprob norm=False: [-4.489243984222412, -5.2529401779174805] manual sum: [-4.489243984222412, -5.2529401779174805] sequence_logprob norm=True: [-1.4964146614074707, -1.750980019569397] manual sum/masksum: [-1.4964146614074707, -1.750980019569397] [stdout] token_logprobs: [[-1.4040248394012451, -2.092271089553833, -0.9929482340812683], [-1.7721467018127441, -1.8669283390045166, -1.6138648986816406]] manual : [[-1.4040248394012451, -2.092271089553833, -0.9929482340812683], [-1.7721467018127441, -1.8669283390045166, -1.6138648986816406]] selective_logprobs: [-4.489243984222412, -5.2529401779174805] masked tl : [[-1.4040248394012451, -2.092271089553833, -0.9929482340812683], [-1.7721467018127441, -1.8669283390045166, -1.6138648986816406]] sequence_logprob norm=False: [-4.489243984222412, -5.2529401779174805] manual sum: [-4.489243984222412, -5.2529401779174805] sequence_logprob norm=True: [-1.4964146614074707, -1.750980019569397] manual sum/masksum: [-1.4964146614074707, -1.750980019569397]
Probe entropy
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(4)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
p=torch.softmax(logits,-1); lp=torch.log_softmax(logits,-1)
ent=-(p*lp).sum(-1)
print('entropy:', R.entropy(logits,mask).item())
print(' masked_mean ent:', ((ent*mask).sum()/mask.sum()).item())
print(' mean ent all:', ent.mean().item())
# normalized_entropy
print('normalized_entropy:', R.normalized_entropy(logits,mask).item())
print(' /logV:', ((ent*mask).sum()/mask.sum()).item()/torch.log(torch.tensor(float(V))).item())
"entropy: 1.233632206916809 masked_mean ent: 1.2336320877075195 mean ent all: 1.187360167503357 normalized_entropy: 0.7664987444877625 /logV: 0.7664986981651061 [stdout] entropy: 1.233632206916809 masked_mean ent: 1.2336320877075195 mean ent all: 1.187360167503357 normalized_entropy: 0.7664987444877625 /logV: 0.7664986981651061
Probe dpo/ipo/bt
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(5)
pc=torch.randn(4,generator=g); pr=torch.randn(4,generator=g)
rc=torch.randn(4,generator=g); rr=torch.randn(4,generator=g)
beta=0.1; ls=0.05
print('dpo_loss:', R.dpo_loss(pc,pr,rc,rr,beta,ls).item())
logits=beta*((pc-pr)-(rc-rr))
l = -F.logsigmoid(logits)*(1-ls) - F.logsigmoid(-logits)*ls
print(' manual mean:', l.mean().item())
# ipo
print('ipo_loss:', R.ipo_loss(pc,pr,rc,rr,beta).item())
h=(pc-pr)-(rc-rr)
ipo=((h - 1/(2*beta))**2).mean()
print(' manual ipo:', ipo.item())
# bradley_terry_logit
print('bt:', R.bradley_terry_logit(rc,rr,beta).tolist())
print(' beta*(rc-rr):', (beta*(rc-rr)).tolist())
"dpo_loss: 0.6925042867660522 manual mean: 0.6925042867660522 ipo_loss: 27.01028823852539 manual ipo: 27.01028823852539 bt: [0.04087548330426216, -0.02329099178314209, -0.12424357235431671, 0.011671716347336769] beta*(rc-rr): [0.04087548330426216, -0.02329099178314209, -0.12424357235431671, 0.011671716347336769] [stdout] dpo_loss: 0.6925042867660522 manual mean: 0.6925042867660522 ipo_loss: 27.01028823852539 manual ipo: 27.01028823852539 bt: [0.04087548330426216, -0.02329099178314209, -0.12424357235431671, 0.011671716347336769] beta*(rc-rr): [0.04087548330426216, -0.02329099178314209, -0.12424357235431671, 0.011671716347336769]
Probe grpo/rloo advantages
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(6)
rewards=torch.randn(6,generator=g)
print('grpo_adv scale=True:', R.grpo_advantages(rewards,3,True).tolist())
r=rewards.view(2,3)
mean=r.mean(-1,keepdim=True); std=r.std(-1,keepdim=True)
print(' (r-mean)/(std+1e-6):', ((r-mean)/(std+1e-6)).view(-1).tolist())
print(' (r-mean)/(std_unb+eps... unbiased True):', ((r-mean)/(r.std(-1,unbiased=True,keepdim=True)+1e-6)).view(-1).tolist())
print('grpo_adv scale=False:', R.grpo_advantages(rewards,3,False).tolist())
print(' (r-mean):', (r-mean).view(-1).tolist())
print('rloo_adv:', R.rloo_advantages(rewards,3).tolist())
# rloo: leave one out baseline
gs=3
rr=rewards.view(2,3)
loo=(rr.sum(-1,keepdim=True)-rr)/(gs-1)
print(' r-loo:', (rr-loo).view(-1).tolist())
print('group_mean_baseline:', R.group_mean_baseline(rewards,3).tolist())
print(' broadcast mean:', mean.expand(2,3).reshape(-1).tolist())
"grpo_adv scale=True: [-1.0755648612976074, -0.2574329674243927, 1.3329977989196777, -0.05098499357700348, -1.1984542608261108, 1.2494393587112427] (r-mean)/(std+1e-6): [-0.8781951665878296, -0.21019317209720612, 1.0883883237838745, -0.041629087179899216, -0.9785341024398804, 1.0201632976531982] (r-mean)/(std_unb+eps... unbiased True): [-0.8781951665878296, -0.21019317209720612, 1.0883883237838745, -0.041629087179899216, -0.9785341024398804, 1.0201632976531982] grpo_adv scale=False: [-1.1578670740127563, -0.27713173627853394, 1.4349987506866455, -0.03435087203979492, -0.8074522614479065, 0.8418031930923462] (r-mean): [-1.1578670740127563, -0.27713173627853394, 1.4349987506866455, -0.03435087203979492, -0.8074522614479065, 0.8418031930923462] rloo_adv: [-1.7368006706237793, -0.4156976342201233, 2.152498245239258, -0.05152636766433716, -1.2111784219741821, 1.262704849243164] r-loo: [-1.7368006706237793, -0.4156976342201233, 2.152498245239258, -0.05152636766433716, -1.2111784219741821, 1.262704849243164] group_mean_baseline: [-1.1578670740127563, -0.27713173627853394, 1.4349987506866455, -0.03435087203979492, -0.8074522614479065, 0.8418031930923462] broadcast mean: [-0.71653151512146, -0.71653151512146, -0.71653151512146, -0.6641029715538025, -0.6641029715538025, -0.6641029715538025] [stdout] grpo_adv scale=True: [-1.0755648612976074, -0.2574329674243927, 1.3329977989196777, -0.05098499357700348, -1.1984542608261108, 1.2494393587112427] (r-mean)/(std+1e-6): [-0.8781951665878296, -0.21019317209720612, 1.0883883237838745, -0.041629087179899216, -0.9785341024398804, 1.0201632976531982] (r-mean)/(std_unb+eps... unbiased True): [-0.8781951665878296, -0.21019317209720612, 1.0883883237838745, -0.041629087179899216, -0.9785341024398804, 1.0201632976531982] grpo_adv scale=False: [-1.1578670740127563, -0.27713173627853394, 1.4349987506866455, -0.03435087203979492, -0.8074522614479065, 0.8418031930923462] (r-mean): [-1.1578670740127563, -0.27713173627853394, 1.4349987506866455, -0.03435087203979492, -0.8074522614479065, 0.8418031930923462] rloo_adv: [-1.7368006706237793, -0.4156976342201233, 2.152498245239258, -0.05152636766433716, -1.2111784219741821, 1.262704849243164] r-loo: [-1.7368006706237793, -0.4156976342201233, 2.152498245239258, -0.05152636766433716, -1.2111784219741821, 1.262704849243164] group_mean_baseline: [-1.1578670740127563, -0.27713173627853394, 1.4349987506866455, -0.03435087203979492, -0.8074522614479065, 0.8418031930923462] broadcast mean: [-0.71653151512146, -0.71653151512146, -0.71653151512146, -0.6641029715538025, -0.6641029715538025, -0.6641029715538025]
Determine std convention for grpo
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(6)
rewards=torch.randn(6,generator=g)
r=rewards.view(2,3)
mean=r.mean(-1,keepdim=True)
for ub in [True,False]:
for eps in [1e-6,1e-8,1e-4]:
std=r.std(-1,unbiased=ub,keepdim=True)
val=((r-mean)/(std+eps)).view(-1)[0].item()
print(f'ub={ub} eps={eps}: {val:.7f}')
print('target:', R.grpo_advantages(rewards,3,True).tolist()[0])
"ub=True eps=1e-06: -0.8781952 ub=True eps=1e-08: -0.8781958 ub=True eps=0.0001: -0.8781292 ub=False eps=1e-06: -1.0755649 ub=False eps=1e-08: -1.0755658 ub=False eps=0.0001: -1.0754659 target: -1.0755648612976074 [stdout] ub=True eps=1e-06: -0.8781952 ub=True eps=1e-08: -0.8781958 ub=True eps=0.0001: -0.8781292 ub=False eps=1e-06: -1.0755649 ub=False eps=1e-08: -1.0755658 ub=False eps=0.0001: -1.0754659 target: -1.0755648612976074
Probe discounted_returns
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(7)
# discounted_returns
rewards=torch.randn(2,4,generator=g)
gamma=0.9
print('discounted_returns:', R.discounted_returns(rewards,gamma).tolist())
# manual
def disc(r,gm):
out=torch.zeros_like(r); acc=torch.zeros(r.shape[0])
for t in reversed(range(r.shape[1])):
acc=r[:,t]+gm*acc; out[:,t]=acc
return out
print(' manual:', disc(rewards,gamma).tolist())
"discounted_returns: [[1.3749160766601562, -0.01920384168624878, 0.6267591118812561, -0.005196809768676758], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] manual: [[0.5153008103370667, 0.7356621623039246, -0.0560879111289978, -1.1143440008163452], [1.4957979917526245, -0.21665799617767334, 0.7535221576690674, 1.2323857545852661]] [stdout] discounted_returns: [[1.3749160766601562, -0.01920384168624878, 0.6267591118812561, -0.005196809768676758], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] manual: [[0.5153008103370667, 0.7356621623039246, -0.0560879111289978, -1.1143440008163452], [1.4957979917526245, -0.21665799617767334, 0.7535221576690674, 1.2323857545852661]]
Determine discounted_returns axis
python3 -c "
import torch, reforacle as R
r=torch.tensor([1.0,0.0,0.0,0.0])
print('1D [1,0,0,0]:', R.discounted_returns(r,0.9).tolist())
r=torch.tensor([0.0,1.0,0.0,0.0])
print('1D [0,1,0,0]:', R.discounted_returns(r,0.9).tolist())
r=torch.tensor([0.0,0.0,0.0,1.0])
print('1D [0,0,0,1]:', R.discounted_returns(r,0.9).tolist())
r2=torch.tensor([[1.0,0.0],[0.0,0.0]])
print('2D col:', R.discounted_returns(r2,0.9).tolist())
"1D [1,0,0,0]: [1.0, 0.0, 0.0, 0.0] 1D [0,1,0,0]: [0.8999999761581421, 1.0, 0.0, 0.0] 1D [0,0,0,1]: [0.7289999127388, 0.809999942779541, 0.8999999761581421, 1.0] 2D col: [[1.0, 0.0], [0.0, 0.0]] [stdout] 1D [1,0,0,0]: [1.0, 0.0, 0.0, 0.0] 1D [0,1,0,0]: [0.8999999761581421, 1.0, 0.0, 0.0] 1D [0,0,0,1]: [0.7289999127388, 0.809999942779541, 0.8999999761581421, 1.0] 2D col: [[1.0, 0.0], [0.0, 0.0]]
Recheck discounted_returns axis on 2x4
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(7)
rewards=torch.randn(2,4,generator=g)
gamma=0.9
tgt=R.discounted_returns(rewards,gamma)
print('rewards:',rewards.tolist())
print('target:',tgt.tolist())
def disc_last(r,gm):
out=torch.zeros_like(r); acc=torch.zeros(r.shape[:-1])
for t in reversed(range(r.shape[-1])):
acc=r[...,t]+gm*acc; out[...,t]=acc
return out
print('last:',disc_last(rewards,gamma).tolist())
def disc_first(r,gm):
out=torch.zeros_like(r); acc=torch.zeros(r.shape[1:])
for t in reversed(range(r.shape[0])):
acc=r[t]+gm*acc; out[t]=acc
return out
print('first:',disc_first(rewards,gamma).tolist())
"rewards: [[-0.1467950940132141, 0.7861412763595581, 0.9468216300010681, -1.1143440008163452], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] target: [[1.3749160766601562, -0.01920384168624878, 0.6267591118812561, -0.005196809768676758], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] last: [[0.5153008103370667, 0.7356621623039246, -0.0560879111289978, -1.1143440008163452], [1.4957979917526245, -0.21665799617767334, 0.7535221576690674, 1.2323857545852661]] first: [[1.3749160766601562, -0.01920384168624878, 0.6267591118812561, -0.005196809768676758], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] [stdout] rewards: [[-0.1467950940132141, 0.7861412763595581, 0.9468216300010681, -1.1143440008163452], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] target: [[1.3749160766601562, -0.01920384168624878, 0.6267591118812561, -0.005196809768676758], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]] last: [[0.5153008103370667, 0.7356621623039246, -0.0560879111289978, -1.1143440008163452], [1.4957979917526245, -0.21665799617767334, 0.7535221576690674, 1.2323857545852661]] first: [[1.3749160766601562, -0.01920384168624878, 0.6267591118812561, -0.005196809768676758], [1.6907901763916016, -0.8948279023170471, -0.3556250333786011, 1.2323857545852661]]
Probe gae and lambda_returns
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(8)
T=4
rewards=torch.randn(T,generator=g)
values=torch.randn(T,generator=g)
next_value=torch.randn(1,generator=g)
gamma,lam=0.9,0.95
adv=R.gae(rewards,values,next_value,gamma,lam)
print('gae:', adv.tolist() if hasattr(adv,'tolist') else adv)
# manual gae
def gae_m(r,v,nv,gm,lm):
T=len(r); adv=torch.zeros(T); last=0.0
for t in reversed(range(T)):
nextv = nv if t==T-1 else v[t+1]
delta=r[t]+gm*nextv-v[t]
last=delta+gm*lm*last
adv[t]=last
return adv
print('manual:', gae_m(rewards,values,next_value,gamma,lam).tolist())
lr=R.lambda_returns(rewards,values,next_value,gamma,lam)
print('lambda_returns:', lr.tolist())
print('gae+values:', (gae_m(rewards,values,next_value,gamma,lam)+values).tolist())
"gae: [-0.4489001929759979, -0.8946236371994019, -0.17952412366867065, -0.8557336926460266] manual: [-0.4489001929759979, -0.8946236371994019, -0.17952412366867065, -0.8557336926460266] lambda_returns: [-1.0883203744888306, -1.5553944110870361, -0.3228384852409363, -0.9600484371185303] gae+values: [-1.0883203744888306, -1.5553944110870361, -0.3228384852409363, -0.9600484371185303] [stdout] gae: [-0.4489001929759979, -0.8946236371994019, -0.17952412366867065, -0.8557336926460266] manual: [-0.4489001929759979, -0.8946236371994019, -0.17952412366867065, -0.8557336926460266] lambda_returns: [-1.0883203744888306, -1.5553944110870361, -0.3228384852409363, -0.9600484371185303] gae+values: [-1.0883203744888306, -1.5553944110870361, -0.3228384852409363, -0.9600484371185303]
Verify gae 2D axis0
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(9)
T,B=4,2
rewards=torch.randn(T,B,generator=g)
values=torch.randn(T,B,generator=g)
next_value=torch.randn(B,generator=g)
gamma,lam=0.9,0.95
adv=R.gae(rewards,values,next_value,gamma,lam)
print('gae 2D:', adv.tolist())
def gae_m(r,v,nv,gm,lm):
T=r.shape[0]; adv=torch.zeros_like(r); last=torch.zeros_like(nv)
for t in reversed(range(T)):
nextv = nv if t==T-1 else v[t+1]
delta=r[t]+gm*nextv-v[t]
last=delta+gm*lm*last
adv[t]=last
return adv
print('manual 2D:', gae_m(rewards,values,next_value,gamma,lam).tolist())
"gae 2D: [[0.9896471500396729, 2.2350449562072754], [2.791564464569092, 1.9602011442184448], [3.5153610706329346, 1.9890869855880737], [1.5665056705474854, 1.5622650384902954]] manual 2D: [[0.9896471500396729, 2.2350449562072754], [2.791564464569092, 1.9602011442184448], [3.5153610706329346, 1.9890869855880737], [1.5665056705474854, 1.5622650384902954]] [stdout] gae 2D: [[0.9896471500396729, 2.2350449562072754], [2.791564464569092, 1.9602011442184448], [3.5153610706329346, 1.9890869855880737], [1.5665056705474854, 1.5622650384902954]] manual 2D: [[0.9896471500396729, 2.2350449562072754], [2.791564464569092, 1.9602011442184448], [3.5153610706329346, 1.9890869855880737], [1.5665056705474854, 1.5622650384902954]]
Probe kl estimators
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(10)
logp=torch.randn(5,generator=g); ref=torch.randn(5,generator=g)
for est in ['k1','k2','k3','kl','mse','abs']:
try:
print('kl_penalty',est,':', R.kl_penalty(logp,ref,est).tolist())
except Exception as e:
print('kl_penalty',est,'ERR',type(e).__name__,str(e)[:40])
d=logp-ref
print(' k1 (logp-ref):', d.tolist())
print(' k3 (exp(-d)-1+d)... ref-logp form:')
lr=ref-logp
print(' k3 = (lr).exp()-1-lr? :', (lr.exp()-1-lr).tolist())
print('reverse_kl:', R.reverse_kl(logp,ref).tolist())
print('symmetric_kl:', R.symmetric_kl(logp,ref).tolist())
"kl_penalty k1 : [-0.2529233992099762, -0.14302963018417358, 0.6559020280838013, -0.03563988208770752, -0.9851673245429993] kl_penalty k2 : [0.031985122710466385, 0.010228737257421017, 0.21510373055934906, 0.0006351005868054926, 0.48527732491493225] kl_penalty k3 : [0.034861207008361816, 0.010734319686889648, 0.17487573623657227, 0.0006426572799682617, 0.6930925846099854] kl_penalty kl ERR ValueError kl kl_penalty mse ERR ValueError mse kl_penalty abs ERR ValueError abs k1 (logp-ref): [-0.2529233992099762, -0.14302963018417358, 0.6559020280838013, -0.03563988208770752, -0.9851673245429993] k3 (exp(-d)-1+d)... ref-logp form: k3 = (lr).exp()-1-lr? : [0.03486117720603943, 0.010734379291534424, 0.17487573623657227, 0.0006426572799682617, 0.6930925250053406] reverse_kl: [0.034861207008361816, 0.010734319686889648, 0.17487573623657227, 0.0006426572799682617, 0.6930925846099854] symmetric_kl: [0.03215599060058594, 0.0102461576461792, 0.2229267954826355, 0.0006351470947265625, 0.5258183479309082] [stdout] kl_penalty k1 : [-0.2529233992099762, -0.14302963018417358, 0.6559020280838013, -0.03563988208770752, -0.9851673245429993] kl_penalty k2 : [0.031985122710466385, 0.010228737257421017, 0.21510373055934906, 0.0006351005868054926, 0.48527732491493225] kl_penalty k3 : [0.034861207008361816, 0.010734319686889648, 0.17487573623657227, 0.0006426572799682617, 0.6930925846099854] kl_penalty kl ERR ValueError kl kl_penalty mse ERR ValueError mse kl_penalty abs ERR ValueError abs k1 (logp-ref): [-0.2529233992099762, -0.14302963018417358, 0.6559020280838013, -0.03563988208770752, -0.9851673245429993] k3 (exp(-d)-1+d)... ref-logp form: k3 = (lr).exp()-1-lr? : [0.03486117720603943, 0.010734379291534424, 0.17487573623657227, 0.0006426572799682617, 0.6930925250053406] reverse_kl: [0.034861207008361816, 0.010734319686889648, 0.17487573623657227, 0.0006426572799682617, 0.6930925846099854] symmetric_kl: [0.03215599060058594, 0.0102461576461792, 0.2229267954826355, 0.0006351470947265625, 0.5258183479309082]
Determine symmetric_kl
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(10)
logp=torch.randn(5,generator=g); ref=torch.randn(5,generator=g)
tgt=R.symmetric_kl(logp,ref)
d=logp-ref
fwd_k3=(d).exp()-1-d # exp(logp-ref)-1-(logp-ref)
rev_k3=(-d).exp()-1+d
print('tgt:',tgt.tolist())
print('fwd_k3+rev_k3:',(fwd_k3+rev_k3).tolist())
print('0.5*(fwd+rev):',(0.5*(fwd_k3+rev_k3)).tolist())
# maybe (p-q)*(logp-logq) style: (exp(logp)-exp(ref))*(logp-ref)
print('(e^lp-e^ref)*d:',((logp.exp()-ref.exp())*d).tolist())
"tgt: [0.03215599060058594, 0.0102461576461792, 0.2229267954826355, 0.0006351470947265625, 0.5258183479309082] fwd_k3+rev_k3: [0.0643119215965271, 0.020492374897003174, 0.445853590965271, 0.001270294189453125, 1.0516365766525269] 0.5*(fwd+rev): [0.03215596079826355, 0.010246187448501587, 0.2229267954826355, 0.0006351470947265625, 0.5258182883262634] (e^lp-e^ref)*d: [0.03989101201295853, 0.007992514409124851, 0.23320287466049194, 0.0003788406611420214, 4.148064613342285] [stdout] tgt: [0.03215599060058594, 0.0102461576461792, 0.2229267954826355, 0.0006351470947265625, 0.5258183479309082] fwd_k3+rev_k3: [0.0643119215965271, 0.020492374897003174, 0.445853590965271, 0.001270294189453125, 1.0516365766525269] 0.5*(fwd+rev): [0.03215596079826355, 0.010246187448501587, 0.2229267954826355, 0.0006351470947265625, 0.5258182883262634] (e^lp-e^ref)*d: [0.03989101201295853, 0.007992514409124851, 0.23320287466049194, 0.0003788406611420214, 4.148064613342285]
Probe importance_ratio, clip_fraction
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(11)
logp=torch.randn(5,generator=g); old=torch.randn(5,generator=g)
print('importance_ratio clip=0.2:', R.importance_ratio(logp,old,0.2).tolist())
r=(logp-old).exp()
print(' clamp exp:', r.clamp(1-0.2,1+0.2).tolist())
print(' raw exp:', r.tolist())
print('importance_ratio clip=None:', R.importance_ratio(logp,old,None).tolist())
# clip_fraction
print('clip_fraction clip=0.2:', R.clip_fraction(logp,old,0.2).item())
print(' frac |r-1|>0.2:', ((r-1).abs()>0.2).float().mean().item())
"importance_ratio clip=0.2: [1.2000000476837158, 1.2000000476837158, 0.800000011920929, 0.800000011920929, 1.2000000476837158] clamp exp: [1.2000000476837158, 1.2000000476837158, 0.800000011920929, 0.800000011920929, 1.2000000476837158] raw exp: [2.738058090209961, 5.472548484802246, 0.3062152862548828, 0.1733027696609497, 1.5582621097564697] importance_ratio clip=None: [2.738058090209961, 5.472548484802246, 0.3062152862548828, 0.1733027696609497, 1.5582621097564697] clip_fraction clip=0.2: 1.0 frac |r-1|>0.2: 1.0 [stdout] importance_ratio clip=0.2: [1.2000000476837158, 1.2000000476837158, 0.800000011920929, 0.800000011920929, 1.2000000476837158] clamp exp: [1.2000000476837158, 1.2000000476837158, 0.800000011920929, 0.800000011920929, 1.2000000476837158] raw exp: [2.738058090209961, 5.472548484802246, 0.3062152862548828, 0.1733027696609497, 1.5582621097564697] importance_ratio clip=None: [2.738058090209961, 5.472548484802246, 0.3062152862548828, 0.1733027696609497, 1.5582621097564697] clip_fraction clip=0.2: 1.0 frac |r-1|>0.2: 1.0
Refine clip_fraction
python3 -c "
import torch, reforacle as R
logp=torch.tensor([0.0,0.1,0.3,-0.05,-0.5])
old=torch.zeros(5)
r=(logp-old).exp()
print('r:',r.tolist())
print('clip_fraction 0.2:', R.clip_fraction(logp,old,0.2).item())
print(' |r-1|>0.2:', ((r-1).abs()>0.2).float().mean().item())
print(' |logp-old|>0.2:', ((logp-old).abs()>0.2).float().mean().item())
"r: [1.0, 1.1051709651947021, 1.3498587608337402, 0.951229453086853, 0.6065306663513184] clip_fraction 0.2: 0.4000000059604645 |r-1|>0.2: 0.4000000059604645 |logp-old|>0.2: 0.4000000059604645 [stdout] r: [1.0, 1.1051709651947021, 1.3498587608337402, 0.951229453086853, 0.6065306663513184] clip_fraction 0.2: 0.4000000059604645 |r-1|>0.2: 0.4000000059604645 |logp-old|>0.2: 0.4000000059604645
Distinguish clip_fraction convention
python3 -c "
import torch, reforacle as R
logp=torch.tensor([0.19,-0.19,0.18,-0.25])
old=torch.zeros(4)
r=(logp-old).exp()
print('r-1:',(r-1).tolist())
print('clip_fraction 0.2:', R.clip_fraction(logp,old,0.2).item())
print(' |r-1|>0.2:', ((r-1).abs()>0.2).float().mean().item())
print(' |logp-old|>0.2:', ((logp-old).abs()>0.2).float().mean().item())
"r-1: [0.2092496156692505, -0.17304086685180664, 0.19721734523773193, -0.22119921445846558] clip_fraction 0.2: 0.5 |r-1|>0.2: 0.5 |logp-old|>0.2: 0.25 [stdout] r-1: [0.2092496156692505, -0.17304086685180664, 0.19721734523773193, -0.22119921445846558] clip_fraction 0.2: 0.5 |r-1|>0.2: 0.5 |logp-old|>0.2: 0.25
Probe clipped_pg_loss
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(12)
logp=torch.randn(2,4,generator=g); old=torch.randn(2,4,generator=g)
adv=torch.randn(2,4,generator=g); mask=(torch.rand(2,4,generator=g)>0.3).float()
cl,ch=0.2,0.2
out=R.clipped_pg_loss(logp,old,adv,mask,cl,ch)
print('clipped_pg_loss:', out.item())
r=(logp-old).exp()
unclipped=-adv*r
clipped=-adv*r.clamp(1-cl,1+ch)
loss=torch.max(unclipped,clipped)
print(' masked_mean max:', ((loss*mask).sum()/mask.sum()).item())
print(' mean(all):', loss.mean().item())
"clipped_pg_loss: -0.06370402872562408 masked_mean max: -0.06370402872562408 mean(all): -0.02508426457643509 [stdout] clipped_pg_loss: -0.06370402872562408 masked_mean max: -0.06370402872562408 mean(all): -0.02508426457643509
Probe value_loss
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(13)
values=torch.randn(2,4,generator=g); old=torch.randn(2,4,generator=g); returns=torch.randn(2,4,generator=g)
clip=0.2
print('value_loss clip=0.2:', R.value_loss(values,old,returns,clip).item())
vc=old+(values-old).clamp(-clip,clip)
l1=(values-returns)**2; l2=(vc-returns)**2
print(' 0.5*mean max:', (0.5*torch.max(l1,l2).mean()).item())
print(' mean max:', (torch.max(l1,l2).mean()).item())
print('value_loss clip=None:', R.value_loss(values,old,returns,None).item())
print(' 0.5*mean l1:', (0.5*l1.mean()).item())
print(' mean l1:', (l1.mean()).item())
"Exit code 1 value_loss clip=0.2: 1.1638469696044922 0.5*mean max: 1.1638469696044922 mean max: 2.3276939392089844 Traceback (most recent call last): File "<string>", line 11, in <module> File "reforacle.py", line 109, in reforacle.value_loss TypeError: bad operand type for unary -: 'NoneType' [error] tool reported failure
Probe whiten
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(14)
v=torch.randn(2,4,generator=g); mask=(torch.rand(2,4,generator=g)>0.3).float()
for sm in [True,False]:
print('whiten sm=',sm,':', R.whiten(v,mask,sm).view(-1).tolist())
mean=v.mean();
print(' (v-mean)/(std+1e-8) unb True:', ((v-v.mean())/(v.std()+1e-8)).view(-1).tolist())
print(' no shift: v/(std):', (v/(v.std()+1e-8)).view(-1).tolist())
"whiten sm= True : [-0.8619763255119324, -0.1025102362036705, -0.5514551997184753, -0.682371973991394, -0.049288056790828705, 0.9539925456047058, 1.8304486274719238, -1.3988155126571655] whiten sm= False : [-1.1473407745361328, -0.38787466287612915, -0.8368196487426758, -0.9677363634109497, -0.33465248346328735, 0.6686280965805054, 1.5450842380523682, -1.6841799020767212] (v-mean)/(std+1e-8) unb True: [-0.7214695811271667, 0.005009348504245281, -0.4244358241558075, -0.5496663451194763, 0.0559198372066021, 1.0156233310699463, 1.854010820388794, -1.2349913120269775] no shift: v/(std): [-1.1474010944366455, -0.42092224955558777, -0.8503674268722534, -0.9755979180335999, -0.37001174688339233, 0.5896916389465332, 1.4280792474746704, -1.660922884941101] [stdout] whiten sm= True : [-0.8619763255119324, -0.1025102362036705, -0.5514551997184753, -0.682371973991394, -0.049288056790828705, 0.9539925456047058, 1.8304486274719238, -1.3988155126571655] whiten sm= False : [-1.1473407745361328, -0.38787466287612915, -0.8368196487426758, -0.9677363634109497, -0.33465248346328735, 0.6686280965805054, 1.5450842380523682, -1.6841799020767212] (v-mean)/(std+1e-8) unb True: [-0.7214695811271667, 0.005009348504245281, -0.4244358241558075, -0.5496663451194763, 0.0559198372066021, 1.0156233310699463, 1.854010820388794, -1.2349913120269775] no shift: v/(std): [-1.1474010944366455, -0.42092224955558777, -0.8503674268722534, -0.9755979180335999, -0.37001174688339233, 0.5896916389465332, 1.4280792474746704, -1.660922884941101]
Determine whiten masked stats
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(14)
v=torch.randn(2,4,generator=g); mask=(torch.rand(2,4,generator=g)>0.3).float()
tgtT=R.whiten(v,mask,True); tgtF=R.whiten(v,mask,False)
mtgtT=R.masked_whiten(v,mask,True)
# masked mean/std
def mmean(v,m): return (v*m).sum()/m.sum()
mean=mmean(v,mask)
n=mask.sum()
var=mmean((v-mean)**2,mask) # biased masked var
std_b=var.sqrt()
# unbiased masked var
var_u=(( (v-mean)**2*mask).sum())/(n-1)
std_u=var_u.sqrt()
print('whiten T:',tgtT.view(-1)[:3].tolist())
print(' masked (v-mean)/(std_b+1e-8):', ((v-mean)/(std_b+1e-8)).view(-1)[:3].tolist())
print(' masked (v-mean)/(std_u+1e-8):', ((v-mean)/(std_u+1e-8)).view(-1)[:3].tolist())
print('masked_whiten T:', mtgtT.view(-1)[:3].tolist())
print('whiten==masked_whiten?', torch.allclose(tgtT,mtgtT))
"whiten T: [-0.8619763255119324, -0.1025102362036705, -0.5514551997184753] masked (v-mean)/(std_b+1e-8): [-0.8619763255119324, -0.1025102362036705, -0.5514551997184753] masked (v-mean)/(std_u+1e-8): [-0.7980349659919739, -0.0949060320854187, -0.5105482339859009] masked_whiten T: [-0.0, -0.1025102362036705, -0.5514551997184753] whiten==masked_whiten? False [stdout] whiten T: [-0.8619763255119324, -0.1025102362036705, -0.5514551997184753] masked (v-mean)/(std_b+1e-8): [-0.8619763255119324, -0.1025102362036705, -0.5514551997184753] masked (v-mean)/(std_u+1e-8): [-0.7980349659919739, -0.0949060320854187, -0.5105482339859009] masked_whiten T: [-0.0, -0.1025102362036705, -0.5514551997184753] whiten==masked_whiten? False
Probe cross_entropy, smoothed_nll, logprob_at_temp
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(15)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
# cross_entropy
print('cross_entropy ignore=-100:', R.cross_entropy(logits,labels,-100).item())
print(' F.cross_entropy flat:', F.cross_entropy(logits.reshape(-1,V),labels.reshape(-1)).item())
lab2=labels.clone(); lab2[0,0]=-100
print('cross_entropy w/ ignore:', R.cross_entropy(logits,lab2,-100).item())
print(' F ce ignore:', F.cross_entropy(logits.reshape(-1,V),lab2.reshape(-1),ignore_index=-100).item())
# smoothed_nll
print('smoothed_nll sm=0.1:', R.smoothed_nll(logits,labels,0.1).item())
lp=F.log_softmax(logits,-1)
nll=-lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
smooth=-lp.mean(-1)
sm=0.1
loss=((1-sm)*nll+sm*smooth)
print(' mean:', loss.mean().item())
# logprob_at_temperature
print('logprob_at_temp temp=2:', R.logprob_at_temperature(logits,labels,2.0).tolist())
lpt=F.log_softmax(logits/2.0,-1)
print(' gather:', lpt.gather(-1,labels.unsqueeze(-1)).squeeze(-1).tolist())
"cross_entropy ignore=-100: 1.5486449003219604 F.cross_entropy flat: 1.54864501953125 cross_entropy w/ ignore: 1.4648430347442627 F ce ignore: 1.4648431539535522 smoothed_nll sm=0.1: 1.6026819944381714 mean: 1.6026819944381714 logprob_at_temp temp=2: [[-1.7649052143096924, -1.3469717502593994, -1.5545666217803955], [-1.076872706413269, -1.9048125743865967, -1.1625369787216187]] gather: [[-1.7649052143096924, -1.3469717502593994, -1.5545666217803955], [-1.076872706413269, -1.9048125743865967, -1.1625369787216187]] [stdout] cross_entropy ignore=-100: 1.5486449003219604 F.cross_entropy flat: 1.54864501953125 cross_entropy w/ ignore: 1.4648430347442627 F ce ignore: 1.4648431539535522 smoothed_nll sm=0.1: 1.6026819944381714 mean: 1.6026819944381714 logprob_at_temp temp=2: [[-1.7649052143096924, -1.3469717502593994, -1.5545666217803955], [-1.076872706413269, -1.9048125743865967, -1.1625369787216187]] gather: [[-1.7649052143096924, -1.3469717502593994, -1.5545666217803955], [-1.076872706413269, -1.9048125743865967, -1.1625369787216187]]
Probe argmax, mode, median, masked_max
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(16)
logits=torch.randn(2,3,5,generator=g)
print('argmax_tokens:', R.argmax_tokens(logits).tolist(), R.argmax_tokens(logits).dtype)
print(' argmax -1:', logits.argmax(-1).tolist())
labels=torch.tensor([[1,2,2,3,2],[0,0,1,1,1]])
print('mode_label:', R.mode_label(labels).tolist(), R.mode_label(labels).dtype)
print(' mode -1:', labels.mode(-1).values.tolist())
rewards=torch.randn(6,generator=g)
print('median_reward:', R.median_reward(rewards).item())
print(' median:', rewards.median().item(), 'quantile.5:', rewards.quantile(0.5).item())
v=torch.randn(2,4,generator=g); mask=(torch.rand(2,4,generator=g)>0.4).float()
print('masked_max:', R.masked_max(v,mask).tolist() if hasattr(R.masked_max(v,mask),'tolist') else R.masked_max(v,mask))
print(' v:',v.tolist(),'mask:',mask.tolist())
"argmax_tokens: [[2, 1, 2], [3, 2, 4]] torch.int64 argmax -1: [[2, 1, 2], [3, 2, 4]] mode_label: [2, 1] torch.int64 mode -1: [2, 1] median_reward: -0.34643951058387756 median: -0.4491531252861023 quantile.5: -0.34643951058387756 masked_max: 2.579069137573242 v: [[1.163411259651184, -1.8571020364761353, 0.47712442278862, 0.33771994709968567], [2.579069137573242, -1.2121334075927734, -1.8726484775543213, 0.5299710631370544]] mask: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]] [stdout] argmax_tokens: [[2, 1, 2], [3, 2, 4]] torch.int64 argmax -1: [[2, 1, 2], [3, 2, 4]] mode_label: [2, 1] torch.int64 mode -1: [2, 1] median_reward: -0.34643951058387756 median: -0.4491531252861023 quantile.5: -0.34643951058387756 masked_max: 2.579069137573242 v: [[1.163411259651184, -1.8571020364761353, 0.47712442278862, 0.33771994709968567], [2.579069137573242, -1.2121334075927734, -1.8726484775543213, 0.5299710631370544]] mask: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]]
Confirm masked_max global, quantile_lower
python3 -c "
import torch, reforacle as R
v=torch.tensor([[5.0,1.0],[2.0,3.0]])
mask=torch.tensor([[0.0,1.0],[1.0,1.0]]) # mask out the 5
print('masked_max:', R.masked_max(v,mask)) # should be 3 not 5
mask2=torch.tensor([[0.0,0.0],[0.0,0.0]])
print('masked_max empty:', R.masked_max(v,mask2))
# quantile_lower
g=torch.Generator().manual_seed(17)
x=torch.randn(7,generator=g)
print('quantile_lower q=0.3:', R.quantile_lower(x,0.3).item())
print(' quantile lower interp:', torch.quantile(x,0.3,interpolation='lower').item())
print(' quantile linear:', torch.quantile(x,0.3).item())
"masked_max: tensor(3.) masked_max empty: tensor(0.) quantile_lower q=0.3: -1.183546781539917 quantile lower interp: -1.183546781539917 quantile linear: -0.7265989184379578 [stdout] masked_max: tensor(3.) masked_max empty: tensor(0.) quantile_lower q=0.3: -1.183546781539917 quantile lower interp: -1.183546781539917 quantile linear: -0.7265989184379578
masked_max neg, advantage_mean_std
python3 -c "
import torch, reforacle as R
v=torch.tensor([-5.0,-2.0,-1.0])
mask=torch.tensor([1.0,0.0,0.0])
print('masked_max neg:', R.masked_max(v,mask)) # -5 if true max, 0 if fallback wrong
# quantile_lower confirmed lower interp
# advantage_mean_std
g=torch.Generator().manual_seed(18)
adv=torch.randn(2,4,generator=g); mask=(torch.rand(2,4,generator=g)>0.3).float()
out=R.advantage_mean_std(adv,mask)
print('advantage_mean_std:', out)
mean=(adv*mask).sum()/mask.sum()
var=((adv-mean)**2*mask).sum()/mask.sum()
print(' masked mean,std_biased:', mean.item(), var.sqrt().item())
var_u=((adv-mean)**2*mask).sum()/(mask.sum()-1)
print(' std_unb:', var_u.sqrt().item())
"masked_max neg: tensor(-5.) advantage_mean_std: tensor([0.2512, 0.4582]) masked mean,std_biased: 0.25119057297706604 0.45815062522888184 std_unb: 0.5018788576126099 [stdout] masked_max neg: tensor(-5.) advantage_mean_std: tensor([0.2512, 0.4582]) masked mean,std_biased: 0.25119057297706604 0.45815062522888184 std_unb: 0.5018788576126099
Probe huber, cumulative_max
python3 -c "
import torch, reforacle as R
# huber_value_loss
g=torch.Generator().manual_seed(19)
v=torch.randn(2,4,generator=g); ret=torch.randn(2,4,generator=g)
import torch.nn.functional as F
print('huber_value_loss delta=1:', R.huber_value_loss(v,ret,1.0).item())
print(' F.huber mean:', F.huber_loss(v,ret,delta=1.0).item())
print(' smooth_l1 mean:', F.smooth_l1_loss(v,ret).item())
print(' 0.5*F.huber:', 0.5*F.huber_loss(v,ret,delta=1.0).item())
# cumulative_max
x=torch.tensor([1.0,3.0,2.0,5.0,4.0])
print('cumulative_max:', R.cumulative_max(x).tolist())
print(' cummax:', x.cummax(0).values.tolist())
x2=torch.randn(2,4,generator=g)
print('cumulative_max 2d:', R.cumulative_max(x2).tolist())
print(' cummax dim -1:', x2.cummax(-1).values.tolist())
print(' cummax dim 0:', x2.cummax(0).values.tolist())
"huber_value_loss delta=1: 1.2547177076339722 F.huber mean: 1.2547177076339722 smooth_l1 mean: 1.2547177076339722 0.5*F.huber: 0.6273588538169861 cumulative_max: [1.0, 1.0, 3.0, 3.0, 5.0] cummax: [1.0, 3.0, 3.0, 5.0, 5.0] cumulative_max 2d: [[-2.3645501136779785, -0.269317626953125, -1.300920009613037, -0.824088990688324], [-2.3645501136779785, -0.269317626953125, -1.300920009613037, -0.824088990688324]] cummax dim -1: [[-2.3645501136779785, -0.269317626953125, -0.269317626953125, -0.269317626953125], [0.049306899309158325, 0.049306899309158325, 0.753285825252533, 0.753285825252533]] cummax dim 0: [[-2.3645501136779785, -0.269317626953125, -1.300920009613037, -0.824088990688324], [0.049306899309158325, -0.269317626953125, 0.753285825252533, 0.1388329714536667]] [stdout] huber_value_loss delta=1: 1.2547177076339722 F.huber mean: 1.2547177076339722 smooth_l1 mean: 1.2547177076339722 0.5*F.huber: 0.6273588538169861 cumulative_max: [1.0, 1.0, 3.0, 3.0, 5.0] cummax: [1.0, 3.0, 3.0, 5.0, 5.0] cumulative_max 2d: [[-2.3645501136779785, -0.269317626953125, -1.300920009613037, -0.824088990688324], [-2.3645501136779785, -0.269317626953125, -1.300920009613037, -0.824088990688324]] cummax dim -1: [[-2.3645501136779785, -0.269317626953125, -0.269317626953125, -0.269317626953125], [0.049306899309158325, 0.049306899309158325, 0.753285825252533, 0.753285825252533]] cummax dim 0: [[-2.3645501136779785, -0.269317626953125, -1.300920009613037, -0.824088990688324], [0.049306899309158325, -0.269317626953125, 0.753285825252533, 0.1388329714536667]]
Confirm cumulative_max dim0, first_nonzero_index
python3 -c "
import torch, reforacle as R
x=torch.tensor([[1.0,9],[3,2],[2,5],[5,1]])
print('cumulative_max 4x2:', R.cumulative_max(x).tolist())
# exclusive cummax along dim0: out[i]=max(x[0..i-1]), out[0]=x[0]
c=x.cummax(0).values
out=torch.cat([x[:1],c[:-1]],0)
print(' exclusive dim0:', out.tolist())
# first_nonzero_index
m=torch.tensor([[0.0,0,1,1],[1,0,0,0],[0,0,0,0]])
print('first_nonzero_index:', R.first_nonzero_index(m).tolist(), R.first_nonzero_index(m).dtype)
print('first_nonzero 1d:', R.first_nonzero_index(torch.tensor([0.0,0,3,0])).tolist())
"cumulative_max 4x2: [[1.0, 9.0], [1.0, 9.0], [3.0, 9.0], [3.0, 9.0]] exclusive dim0: [[1.0, 9.0], [1.0, 9.0], [3.0, 9.0], [3.0, 9.0]] first_nonzero_index: [2, 0, -1] torch.int64 first_nonzero 1d: 2 [stdout] cumulative_max 4x2: [[1.0, 9.0], [1.0, 9.0], [3.0, 9.0], [3.0, 9.0]] exclusive dim0: [[1.0, 9.0], [1.0, 9.0], [3.0, 9.0], [3.0, 9.0]] first_nonzero_index: [2, 0, -1] torch.int64 first_nonzero 1d: 2
Probe pad_mask, bucketize
python3 -c "
import torch, reforacle as R
lengths=torch.tensor([2,0,4,3])
print('pad_mask_from_lengths max=4:', R.pad_mask_from_lengths(lengths,4).tolist(), R.pad_mask_from_lengths(lengths,4).dtype)
# bucketize
rewards=torch.tensor([-1.0,0.5,1.5,2.5,0.0,1.0])
b=torch.tensor([0.0,1.0,2.0])
print('bucketize:', R.bucketize_reward(rewards,b).tolist(), R.bucketize_reward(rewards,b).dtype)
print(' bucketize right=F:', torch.bucketize(rewards,b,right=False).tolist())
print(' bucketize right=T:', torch.bucketize(rewards,b,right=True).tolist())
"pad_mask_from_lengths max=4: [[1.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0]] torch.float32 bucketize: [0, 1, 2, 3, 1, 2] torch.int64 bucketize right=F: [0, 1, 2, 3, 0, 1] bucketize right=T: [0, 1, 2, 3, 1, 2] [stdout] pad_mask_from_lengths max=4: [[1.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0]] torch.float32 bucketize: [0, 1, 2, 3, 1, 2] torch.int64 bucketize right=F: [0, 1, 2, 3, 0, 1] bucketize right=T: [0, 1, 2, 3, 1, 2]
Probe top_p, top_k mask
python3 -c "
import torch, reforacle as R
probs=torch.tensor([[0.1,0.2,0.4,0.25,0.05]])
print('top_p_mask p=0.7:', R.top_p_mask(probs,0.7).tolist(), R.top_p_mask(probs,0.7).dtype)
# top_k_mask
logits=torch.tensor([[1.0,3.0,2.0,5.0,4.0]])
print('top_k_mask k=2:', R.top_k_mask(logits,2).tolist(), R.top_k_mask(logits,2).dtype)
"top_p_mask p=0.7: [[0.0, 0.23529411852359772, 0.47058823704719543, 0.29411762952804565, 0.0]] torch.float32 top_k_mask k=2: [[False, False, False, True, True]] torch.bool [stdout] top_p_mask p=0.7: [[0.0, 0.23529411852359772, 0.47058823704719543, 0.29411762952804565, 0.0]] torch.float32 top_k_mask k=2: [[False, False, False, True, True]] torch.bool
Probe dpo_sequence_loss
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(20)
B,T,V=2,3,5
def mk(): return torch.randn(B,T,V,generator=g)
pc,pr,rc,rr=mk(),mk(),mk(),mk()
cl=torch.randint(0,V,(B,T),generator=g); rl=torch.randint(0,V,(B,T),generator=g)
cm=(torch.rand(B,T,generator=g)>0.3).float(); rm=(torch.rand(B,T,generator=g)>0.3).float()
beta,ls=0.1,0.0
out=R.dpo_sequence_loss(pc,pr,rc,rl,rl if False else rl,cm,rm,beta,ls)
" 2>&1 | head -5
echo "---signature check---"
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(20)
B,T,V=2,3,5
def mk(): return torch.randn(B,T,V,generator=g)
pc_l,pr_l,rc_l,rr_l=mk(),mk(),mk(),mk()
cl=torch.randint(0,V,(B,T),generator=g); rl=torch.randint(0,V,(B,T),generator=g)
cm=(torch.rand(B,T,generator=g)>0.3).float(); rm=(torch.rand(B,T,generator=g)>0.3).float()
beta,ls=0.1,0.0
out=R.dpo_sequence_loss(pc_l,pr_l,rc_l,rr_l,cl,rl,cm,rm,beta,ls)
print('dpo_seq:', out.item())
def seqlp(logits,labels,mask):
lp=F.log_softmax(logits,-1)
tl=lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
return (tl*mask).sum(-1)
pc=seqlp(pc_l,cl,cm); pr=seqlp(pr_l,rl,rm); rc=seqlp(rc_l,cl,cm); rr=seqlp(rr_l,rl,rm)
logits=beta*((pc-pr)-(rc-rr))
loss=(-(1-ls)*F.logsigmoid(logits)-ls*F.logsigmoid(-logits)).mean()
print(' manual:', loss.item())
"Traceback (most recent call last): File "<string>", line 11, in <module> File "reforacle.py", line 121, in reforacle.dpo_sequence_loss TypeError: dpo_sequence_loss() takes exactly 10 positional arguments (9 given) ---signature check--- dpo_seq: 0.6778793334960938 manual: 0.6778793334960938 [stdout] Traceback (most recent call last): File "<string>", line 11, in <module> File "reforacle.py", line 121, in reforacle.dpo_sequence_loss TypeError: dpo_sequence_loss() takes exactly 10 positional arguments (9 given) ---signature check--- dpo_seq: 0.6778793334960938 manual: 0.6778793334960938
Probe grpo_objective
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(21)
B,T,V=4,3,5
logits=torch.randn(B,T,V,generator=g)
old=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
cm=(torch.rand(B,T,generator=g)>0.3).float()
rewards=torch.randn(B,generator=g)
gs=2; beta=0.05; cl,ch=0.2,0.2
out=R.grpo_objective(logits,old,ref,labels,cm,rewards,gs,beta,cl,ch,True,'k3')
print('grpo_objective:', out.item())
def tlp(lg):
lp=F.log_softmax(lg,-1); return lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
logp=tlp(logits); oldp=tlp(old); refp=tlp(ref)
# advantages
r=rewards.view(B//gs,gs)
mean=r.mean(-1,keepdim=True); std=r.std(-1,unbiased=False,keepdim=True)
adv=((r-mean)/(std+1e-6)).view(-1) # per sequence
adv_t=adv.unsqueeze(-1) # (B,1)
ratio=(logp-oldp).exp()
pg=torch.max(-adv_t*ratio, -adv_t*ratio.clamp(1-cl,1+ch))
kl=(refp-logp).exp()-1-(refp-logp) # k3
per_tok=pg+beta*kl
# masked mean over all
loss=(per_tok*cm).sum()/cm.sum()
print(' hyp masked_mean all:', loss.item())
# alt: per-sequence mean then mean
seq=(per_tok*cm).sum(-1)/cm.sum(-1)
print(' hyp per-seq then mean:', seq.mean().item())
"grpo_objective: 1.1932880878448486 hyp masked_mean all: 1.1932880878448486 hyp per-seq then mean: 0.7943323850631714 [stdout] grpo_objective: 1.1932880878448486 hyp masked_mean all: 1.1932880878448486 hyp per-seq then mean: 0.7943323850631714
Verify grpo kl dispatch
python3 -c " import torch, reforacle as R import torch.nn.functional as F g=torch.Generator().manual_seed(21) B,T,V=4,3,5 logits=torch.randn(B,T,V,generator=g); old=torch.randn(B,T,V,generator=g); ref=torch.randn(B,T,V,generator=g) labels=torch.randint(0,V,(B,T),generator=g); cm=(torch.rand(B,T,generator=g)>0.3).float() rewards=torch.randn(B,generator=g); gs=2; beta=0.05; cl,ch=0.2,0.2 def tlp(lg): lp=F.log_softmax(lg,-1); return lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1) logp=tlp(logits); oldp=tlp(old); refp=tlp(ref) r=rewards.view(B//gs,gs); mean=r.mean(-1,keepdim=True); std=r.std(-1,unbiased=False,keepdim=True) adv=((r-mean)/(std+1e-6)).view(-1).unsqueeze(-1) ratio=(logp-oldp).exp() pg=torch.max(-adv*ratio, -adv*ratio.clamp(1-cl,1+ch)) def klf(est): d=logp-refp if est=='k1': return d if est=='k2': return 0.5*d*d if est=='k3': return (-d).exp()-1+d for est in ['k1','k2','k3']: out=R.grpo_objective(logits,old,ref,labels,cm,rewards,gs,beta,cl,ch,True,est) loss=((pg+beta*klf(est))*cm).sum()/cm.sum() print(est, out.item(), loss.item()) "
k1 1.173000693321228 1.173000693321228 k2 1.1944141387939453 1.1944141387939453 k3 1.1932880878448486 1.1932880878448486 [stdout] k1 1.173000693321228 1.173000693321228 k2 1.1944141387939453 1.1944141387939453 k3 1.1932880878448486 1.1932880878448486
Probe ppo_objective
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(22)
T=5
rewards=torch.randn(T,generator=g); values=torch.randn(T,generator=g); old_values=torch.randn(T,generator=g)
logp=torch.randn(T,generator=g); old_logp=torch.randn(T,generator=g); next_value=torch.randn(1,generator=g)
gamma,lam=0.99,0.95; cl,ch=0.2,0.2; vf_clip=0.2; vf_coef=0.5
out=R.ppo_objective(rewards,values,old_values,logp,old_logp,next_value,gamma,lam,cl,ch,vf_clip,vf_coef)
print('ppo_objective:', out.item())
def gae_m(r,v,nv,gm,lm):
T=r.shape[0]; adv=torch.zeros_like(r); last=torch.zeros(())
for t in reversed(range(T)):
nextv=nv if t==T-1 else v[t+1]
delta=r[t]+gm*nextv-v[t]; last=delta+gm*lm*last; adv[t]=last
return adv
adv=gae_m(rewards,values,next_value,gamma,lam)
if adv.dim()>1: pass
adv=adv.squeeze()
returns=adv+values
ratio=(logp-old_logp).exp()
pg=torch.max(-adv*ratio,-adv*ratio.clamp(1-cl,1+ch)).mean()
vc=old_values+(values-old_values).clamp(-vf_clip,vf_clip)
vf=0.5*torch.max((values-returns)**2,(vc-returns)**2).mean()
print(' no whiten:', (pg+vf_coef*vf).item())
# whitened adv (unmasked whiten biased)
def whiten(v):
m=v.mean(); var=((v-m)**2).mean(); return (v-m)/(var.sqrt()+1e-8)
advw=whiten(adv)
pgw=torch.max(-advw*ratio,-advw*ratio.clamp(1-cl,1+ch)).mean()
print(' whiten adv:', (pgw+vf_coef*vf).item())
"ppo_objective: 3.4117684364318848 no whiten: 9.573343276977539 whiten adv: 3.4117684364318848 [stdout] ppo_objective: 3.4117684364318848 no whiten: 9.573343276977539 whiten adv: 3.4117684364318848
Verify ppo 2D global whiten
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(23)
T,B=5,3
rewards=torch.randn(T,B,generator=g); values=torch.randn(T,B,generator=g); old_values=torch.randn(T,B,generator=g)
logp=torch.randn(T,B,generator=g); old_logp=torch.randn(T,B,generator=g); next_value=torch.randn(B,generator=g)
gamma,lam=0.99,0.95; cl,ch=0.2,0.2; vf_clip=0.2; vf_coef=0.5
out=R.ppo_objective(rewards,values,old_values,logp,old_logp,next_value,gamma,lam,cl,ch,vf_clip,vf_coef)
print('ppo 2D:', out.item())
def gae_m(r,v,nv,gm,lm):
T=r.shape[0]; adv=torch.zeros_like(r); last=torch.zeros_like(v[0])
for t in reversed(range(T)):
nextv=nv if t==T-1 else v[t+1]
delta=r[t]+gm*nextv-v[t]; last=delta+gm*lm*last; adv[t]=last
return adv
adv=gae_m(rewards,values,next_value,gamma,lam)
returns=adv+values
ratio=(logp-old_logp).exp()
def whiten(v): m=v.mean(); var=((v-m)**2).mean(); return (v-m)/(var.sqrt()+1e-8)
advw=whiten(adv)
pgw=torch.max(-advw*ratio,-advw*ratio.clamp(1-cl,1+ch)).mean()
vc=old_values+(values-old_values).clamp(-vf_clip,vf_clip)
vf=0.5*torch.max((values-returns)**2,(vc-returns)**2).mean()
print(' hyp:', (pgw+vf_coef*vf).item())
"ppo 2D: 2.994100570678711 hyp: 2.994100570678711 [stdout] ppo 2D: 2.994100570678711 hyp: 2.994100570678711
Probe rloo_objective
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(24)
B,T,V=4,3,5
logits=torch.randn(B,T,V,generator=g); old=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g); mask=(torch.rand(B,T,generator=g)>0.3).float()
rewards=torch.randn(B,generator=g); gs=2; cl,ch=0.2,0.2
out=R.rloo_objective(logits,old,labels,mask,rewards,gs,cl,ch)
print('rloo_objective:', out.item())
def tlp(lg): lp=F.log_softmax(lg,-1); return lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
logp=tlp(logits); oldp=tlp(old)
rr=rewards.view(B//gs,gs); loo=(rr.sum(-1,keepdim=True)-rr)/(gs-1); adv=(rr-loo).view(-1).unsqueeze(-1)
ratio=(logp-oldp).exp()
pg=torch.max(-adv*ratio,-adv*ratio.clamp(1-cl,1+ch))
print(' masked_mean:', ((pg*mask).sum()/mask.sum()).item())
"rloo_objective: 0.9098412990570068 masked_mean: 0.9098412990570068 [stdout] rloo_objective: 0.9098412990570068 masked_mean: 0.9098412990570068
Probe reverse_kl_objective
python3 -c "
import torch, reforacle as R
import torch.nn.functional as F
g=torch.Generator().manual_seed(25)
B,T,V=3,4,5
logits=torch.randn(B,T,V,generator=g); ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g); mask=(torch.rand(B,T,generator=g)>0.3).float()
advantages=torch.randn(B,generator=g); beta=0.1
out=R.reverse_kl_objective(logits,ref,labels,mask,advantages,beta)
print('reverse_kl_objective:', out.item())
def tlp(lg): lp=F.log_softmax(lg,-1); return lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
logp=tlp(logits); refp=tlp(ref)
adv=advantages.unsqueeze(-1)
kl=(refp-logp).exp()-1-(refp-logp) # reverse_kl k3
# pg term: -adv*logp ?
pg=-adv*logp
per=pg+beta*kl
print(' -adv*logp + beta*kl masked_mean:', ((per*mask).sum()/mask.sum()).item())
# maybe adv shape per token? advantages given as (B,)
"reverse_kl_objective: 1.1568137407302856 -adv*logp + beta*kl masked_mean: 1.1568137407302856 [stdout] reverse_kl_objective: 1.1568137407302856 -adv*logp + beta*kl masked_mean: 1.1568137407302856
Check masked_mean empty/denominator
python3 probe.py 2>&1 | tail -6
echo "---"
python3 -c "
import torch, reforacle as R
# masked_mean denominator: use a mask sum, check clamp min
v=torch.tensor([2.0,4.0]); m=torch.tensor([1.0,0.0])
print('mm partial:', R.masked_mean(v,m).item()) # 2/1=2
# masked_sum axis None
print('ms none:', R.masked_sum(torch.tensor([1.0,2,3]),torch.tensor([1.0,1,0])).item())
# does masked_mean clamp denom to 1 or eps? empty numerator 0 -> can't tell, but test all-masked with axis
v=torch.tensor([[5.0,5.0],[1.0,1.0]]); m=torch.tensor([[0.0,0.0],[1.0,1.0]])
print('mm axis1 empty row:', R.masked_mean(v,m,axis=1).tolist())
"mask = [1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0] reference masked_mean -> -0.07950486242771149 reference masked_mean (empty mask) -> 0.0 Probe any reference function this way to discover its exact numerical convention, then implement it in /app/objectives.py to match. --- mm partial: 2.0 ms none: 3.0 mm axis1 empty row: [0.0, 1.0] [stdout] mask = [1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0] reference masked_mean -> -0.07950486242771149 reference masked_mean (empty mask) -> 0.0 Probe any reference function this way to discover its exact numerical convention, then implement it in /app/objectives.py to match. --- mm partial: 2.0 ms none: 3.0 mm axis1 empty row: [0.0, 1.0]
from __future__ import annotations
import torch
import torch.nn.functional as F
def masked_mean(values, mask, axis=None):
num = (values * mask).sum(axis) if axis is not None else (values * mask).sum()
den = mask.sum(axis) if axis is not None else mask.sum()
return num / den.clamp(min=1e-8)
def masked_sum(values, mask, axis=None):
if axis is not None:
return (values * mask).sum(axis)
return (values * mask).sum()
def logsumexp(x, axis):
return torch.logsumexp(x, axis)
def log_softmax(x, axis):
return torch.log_softmax(x, axis)
def _token_logprobs(logits, labels):
lp = torch.log_softmax(logits, -1)
return lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
def selective_logprobs(logits, labels, mask):
tl = _token_logprobs(logits, labels)
return (tl * mask).sum(-1)
def token_logprobs(logits, labels):
return _token_logprobs(logits, labels)
def sequence_logprob(logits, labels, mask, length_normalize):
tl = _token_logprobs(logits, labels)
s = (tl * mask).sum(-1)
if length_normalize:
return s / mask.sum(-1).clamp(min=1e-8)
return s
def entropy(logits, mask):
lp = torch.log_softmax(logits, -1)
p = lp.exp()
ent = -(p * lp).sum(-1)
return (ent * mask).sum() / mask.sum().clamp(min=1e-8)
def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
logits = beta * ((pc - pr) - (rc - rr))
return (-(1 - label_smoothing) * F.logsigmoid(logits)
- label_smoothing * F.logsigmoid(-logits)).mean()
def ipo_loss(pc, pr, rc, rr, beta):
h = (pc - pr) - (rc - rr)
return ((h - 1.0 / (2.0 * beta)) ** 2).mean()
def grpo_advantages(rewards, group_size, scale_by_std):
r = rewards.view(-1, group_size)
mean = r.mean(-1, keepdim=True)
adv = r - mean
if scale_by_std:
std = r.std(-1, unbiased=False, keepdim=True)
adv = adv / (std + 1e-6)
return adv.reshape(rewards.shape)
def gae(rewards, values, next_value, gamma, lam):
T = rewards.shape[0]
adv = torch.zeros_like(rewards)
last = torch.zeros_like(values[0]) if values.dim() > 1 else torch.zeros(())
for t in reversed(range(T)):
nextv = next_value if t == T - 1 else values[t + 1]
delta = rewards[t] + gamma * nextv - values[t]
last = delta + gamma * lam * last
adv[t] = last
return adv
def kl_penalty(logp, ref_logp, estimator):
d = logp - ref_logp
if estimator == 'k1':
return d
if estimator == 'k2':
return 0.5 * d * d
if estimator == 'k3':
return (-d).exp() - 1 + d
raise ValueError(estimator)
def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
ratio = (logp - old_logp).exp()
unclipped = -advantages * ratio
clipped = -advantages * ratio.clamp(1 - clip_low, 1 + clip_high)
loss = torch.max(unclipped, clipped)
return (loss * mask).sum() / mask.sum().clamp(min=1e-8)
def value_loss(values, old_values, returns, clip):
l1 = (values - returns) ** 2
if clip is not None:
vc = old_values + (values - old_values).clamp(-clip, clip)
l2 = (vc - returns) ** 2
return 0.5 * torch.max(l1, l2).mean()
return 0.5 * l1.mean()
def _whiten(values, mask, shift_mean):
if mask is None:
mean = values.mean()
var = ((values - mean) ** 2).mean()
else:
den = mask.sum().clamp(min=1e-8)
mean = (values * mask).sum() / den
var = ((values - mean) ** 2 * mask).sum() / den
std = var.sqrt()
out = values / (std + 1e-8) if not shift_mean else (values - mean) / (std + 1e-8)
return out
def whiten(values, mask, shift_mean):
return _whiten(values, mask, shift_mean)
def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
chosen_labels, rejected_labels, chosen_mask, rejected_mask,
beta, label_smoothing):
def seqlp(logits, labels, mask):
tl = _token_logprobs(logits, labels)
return (tl * mask).sum(-1)
pc = seqlp(pc_logits, chosen_labels, chosen_mask)
pr = seqlp(pr_logits, rejected_labels, rejected_mask)
rc = seqlp(rc_logits, chosen_labels, chosen_mask)
rr = seqlp(rr_logits, rejected_labels, rejected_mask)
return dpo_loss(pc, pr, rc, rr, beta, label_smoothing)
def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
rewards, group_size, beta, clip_low, clip_high, scale_by_std,
kl_estimator):
logp = _token_logprobs(logits, labels)
oldp = _token_logprobs(old_logits, labels)
refp = _token_logprobs(ref_logits, labels)
adv = grpo_advantages(rewards, group_size, scale_by_std).reshape(-1).unsqueeze(-1)
ratio = (logp - oldp).exp()
pg = torch.max(-adv * ratio, -adv * ratio.clamp(1 - clip_low, 1 + clip_high))
kl = kl_penalty(logp, refp, kl_estimator)
per_tok = pg + beta * kl
return (per_tok * completion_mask).sum() / completion_mask.sum().clamp(min=1e-8)
def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
adv = gae(rewards, values, next_value, gamma, lam)
returns = adv + values
advw = _whiten(adv, None, True)
ratio = (logp - old_logp).exp()
pg = torch.max(-advw * ratio, -advw * ratio.clamp(1 - clip_low, 1 + clip_high)).mean()
vf = value_loss(values, old_values, returns, vf_clip)
return pg + vf_coef * vf
def rloo_advantages(rewards, group_size):
r = rewards.view(-1, group_size)
loo = (r.sum(-1, keepdim=True) - r) / (group_size - 1)
return (r - loo).reshape(rewards.shape)
def reverse_kl(logp, ref_logp):
d = ref_logp - logp
return d.exp() - 1 - d
def importance_ratio(logp, old_logp, clip):
r = (logp - old_logp).exp()
if clip is not None:
return r.clamp(1 - clip, 1 + clip)
return r
def discounted_returns(rewards, gamma):
T = rewards.shape[0]
out = torch.zeros_like(rewards)
acc = torch.zeros_like(rewards[0]) if rewards.dim() > 1 else torch.zeros(())
for t in reversed(range(T)):
acc = rewards[t] + gamma * acc
out[t] = acc
return out
def normalize(x, eps):
return (x - x.mean()) / (x.std() + eps)
def top_p_mask(probs, p):
sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
cumsum = sorted_probs.cumsum(-1)
remove = cumsum > p
remove[..., 1:] = remove[..., :-1].clone()
remove[..., 0] = False
sorted_probs = sorted_probs.masked_fill(remove, 0.0)
out = torch.zeros_like(probs)
out.scatter_(-1, sorted_idx, sorted_probs)
out = out / out.sum(-1, keepdim=True)
return out
def smoothed_nll(logits, labels, smoothing):
lp = torch.log_softmax(logits, -1)
nll = -lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
smooth = -lp.mean(-1)
return ((1 - smoothing) * nll + smoothing * smooth).mean()
def bradley_terry_logit(chosen_reward, rejected_reward, beta):
return beta * (chosen_reward - rejected_reward)
def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
logp = _token_logprobs(logits, labels)
oldp = _token_logprobs(old_logits, labels)
adv = rloo_advantages(rewards, group_size).reshape(-1).unsqueeze(-1)
ratio = (logp - oldp).exp()
pg = torch.max(-adv * ratio, -adv * ratio.clamp(1 - clip_low, 1 + clip_high))
return (pg * mask).sum() / mask.sum().clamp(min=1e-8)
def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
logp = _token_logprobs(logits, labels)
refp = _token_logprobs(ref_logits, labels)
adv = advantages.unsqueeze(-1)
kl = reverse_kl(logp, refp)
per_tok = -adv * logp + beta * kl
return (per_tok * mask).sum() / mask.sum().clamp(min=1e-8)
def cross_entropy(logits, labels, ignore_index):
V = logits.shape[-1]
return F.cross_entropy(logits.reshape(-1, V), labels.reshape(-1),
ignore_index=ignore_index)
def top_k_mask(logits, k):
kth = torch.topk(logits, k, dim=-1).values[..., -1:]
return logits >= kth
def group_mean_baseline(rewards, group_size):
r = rewards.view(-1, group_size)
mean = r.mean(-1, keepdim=True)
return (r - mean).reshape(rewards.shape)
def lambda_returns(rewards, values, next_value, gamma, lam):
return gae(rewards, values, next_value, gamma, lam) + values
def symmetric_kl(logp, ref_logp):
d = logp - ref_logp
fwd = d.exp() - 1 - d
rev = (-d).exp() - 1 + d
return 0.5 * (fwd + rev)
def huber_value_loss(values, returns, delta):
return F.huber_loss(values, returns, delta=delta)
def normalized_entropy(logits, mask):
V = logits.shape[-1]
return entropy(logits, mask) / torch.log(torch.tensor(float(V)))
def clip_fraction(logp, old_logp, clip):
r = (logp - old_logp).exp()
return ((r - 1).abs() > clip).float().mean()
def masked_whiten(values, mask, shift_mean):
return _whiten(values, mask, shift_mean) * mask
def logprob_at_temperature(logits, labels, temperature):
lp = torch.log_softmax(logits / temperature, -1)
return lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
def advantage_mean_std(advantages, mask):
den = mask.sum().clamp(min=1e-8)
mean = (advantages * mask).sum() / den
var = ((advantages - mean) ** 2 * mask).sum() / den
return torch.stack([mean, var.sqrt()])
def argmax_tokens(logits):
return logits.argmax(-1)
def mode_label(labels):
return labels.mode(-1).values
def median_reward(rewards):
return rewards.quantile(0.5)
def masked_max(values, mask):
if mask.sum() == 0:
return torch.zeros((), dtype=values.dtype, device=values.device)
return values.masked_fill(mask == 0, float('-inf')).max()
def quantile_lower(x, q):
return torch.quantile(x, q, interpolation='lower')
def pad_mask_from_lengths(lengths, max_len):
ar = torch.arange(max_len, device=lengths.device)
return (ar.unsqueeze(0) < lengths.unsqueeze(-1)).float()
def first_nonzero_index(mask):
nonzero = mask != 0
idx = nonzero.float().argmax(-1)
has = nonzero.any(-1)
return torch.where(has, idx, torch.full_like(idx, -1))
def cumulative_max(x):
c = x.cummax(0).values
return torch.cat([x[:1], c[:-1]], 0)
def bucketize_reward(rewards, boundaries):
return torch.bucketize(rewards, boundaries, right=True)
Comprehensive test all functions
python3 -c "
import torch, reforacle as R, objectives as O
import random
def close(a,b,atol=1e-6,rtol=1e-5):
a=torch.as_tensor(a); b=torch.as_tensor(b)
if a.shape!=b.shape: return False, f'shape {a.shape} vs {b.shape}'
ok=torch.allclose(a.float(),b.float(),atol=atol,rtol=rtol)
return ok, '' if ok else f'max diff {(a.float()-b.float()).abs().max().item()}'
fails={}
def chk(name,args,atol=1e-6,rtol=1e-5):
try:
r=getattr(R,name)(*args); o=getattr(O,name)(*args)
except Exception as e:
fails.setdefault(name,[]).append(f'EXC {type(e).__name__}:{e}'); return
ok,msg=close(r,o,atol,rtol)
if not ok: fails.setdefault(name,[]).append(msg)
for seed in range(20):
g=torch.Generator().manual_seed(seed)
def rn(*s): return torch.randn(*s,generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
def rm(*s): return (torch.rand(*s,generator=g)>0.3).float()
B,T,V=random.Random(seed).choice([(2,3,5),(4,3,6),(3,5,4),(6,2,7)])
v=rn(B,T); m=rm(B,T)
chk('masked_mean',(v,m)); chk('masked_mean',(v,m,0)); chk('masked_mean',(v,m,1))
chk('masked_sum',(v,m)); chk('masked_sum',(v,m,1))
chk('logsumexp',(rn(B,T),1)); chk('log_softmax',(rn(B,T),1))
logits=rn(B,T,V); labels=ri(V,B,T)
chk('selective_logprobs',(logits,labels,m)); chk('token_logprobs',(logits,labels))
chk('sequence_logprob',(logits,labels,m,True)); chk('sequence_logprob',(logits,labels,m,False))
chk('entropy',(logits,m)); chk('normalized_entropy',(logits,m))
pc,pr,rc,rr=rn(B),rn(B),rn(B),rn(B)
chk('dpo_loss',(pc,pr,rc,rr,0.1,0.0)); chk('dpo_loss',(pc,pr,rc,rr,0.5,0.1))
chk('ipo_loss',(pc,pr,rc,rr,0.3)); chk('bradley_terry_logit',(rc,rr,0.2))
gs=random.Random(seed).choice([2,3])
rew=rn(gs*3)
chk('grpo_advantages',(rew,gs,True)); chk('grpo_advantages',(rew,gs,False))
chk('rloo_advantages',(rew,gs)); chk('group_mean_baseline',(rew,gs))
Tt=random.Random(seed).choice([4,5,6])
chk('gae',(rn(Tt),rn(Tt),rn(1),0.99,0.95),1e-5,1e-4)
chk('gae',(rn(Tt,2),rn(Tt,2),rn(2),0.9,0.9),1e-5,1e-4)
chk('lambda_returns',(rn(Tt),rn(Tt),rn(1),0.95,0.9),1e-5,1e-4)
chk('discounted_returns',(rn(Tt,2),0.9),1e-5,1e-4); chk('discounted_returns',(rn(Tt),0.95),1e-5,1e-4)
lp,ref=rn(B,T),rn(B,T)
for est in ['k1','k2','k3']: chk('kl_penalty',(lp,ref,est))
chk('reverse_kl',(lp,ref)); chk('symmetric_kl',(lp,ref))
chk('clipped_pg_loss',(lp,rn(B,T),rn(B,T),m,0.2,0.2))
chk('value_loss',(rn(B,T),rn(B,T),rn(B,T),0.2))
chk('whiten',(v,m,True)); chk('whiten',(v,m,False))
chk('masked_whiten',(v,m,True)); chk('masked_whiten',(v,m,False))
chk('importance_ratio',(lp,rn(B,T),0.2)); chk('importance_ratio',(lp,rn(B,T),None))
chk('normalize',(rn(B,T),1e-5))
chk('smoothed_nll',(logits,labels,0.1))
chk('cross_entropy',(logits,labels,-100))
lab2=labels.clone(); lab2[0,0]=-100; chk('cross_entropy',(logits,lab2,-100))
chk('logprob_at_temperature',(logits,labels,2.0)); chk('logprob_at_temperature',(logits,labels,0.7))
chk('clip_fraction',(lp,rn(B,T),0.2))
chk('advantage_mean_std',(rn(B,T),m))
chk('argmax_tokens',(logits)); chk('mode_label',(labels))
chk('median_reward',(rn(9))); chk('masked_max',(v,m))
chk('quantile_lower',(rn(11),0.3)); chk('quantile_lower',(rn(11),0.75))
chk('huber_value_loss',(rn(B,T),rn(B,T),1.0)); chk('huber_value_loss',(rn(B,T),rn(B,T),0.5))
chk('cumulative_max',(rn(6))); chk('cumulative_max',(rn(5,3)))
lens=ri(T+1,B); chk('pad_mask_from_lengths',(lens,T))
chk('first_nonzero_index',(rm(B,T)))
probs=torch.softmax(rn(B,V),-1); chk('top_p_mask',(probs,0.7)); chk('top_p_mask',(probs,0.9))
chk('top_k_mask',(rn(B,V),2)); chk('top_k_mask',(rn(B,V),3))
bnd,_=torch.sort(rn(3)); chk('bucketize_reward',(rn(8),bnd))
print('FAILS:', {k:v[:2] for k,v in fails.items()} if fails else 'NONE')
print('total funcs tested with fails:', len(fails))
"FAILS: {'whiten': ['max diff 0.03969073295593262', 'max diff 0.03177332878112793'], 'masked_whiten': ['max diff 0.1814074069261551', 'max diff 0.09044408798217773'], 'argmax_tokens': ['EXC TypeError:argmax_tokens() takes exactly 1 positional argument (6 given)', 'EXC TypeError:argmax_tokens() takes exactly 1 positional argument (4 given)'], 'mode_label': ['EXC TypeError:mode_label() takes exactly 1 positional argument (6 given)', 'EXC TypeError:mode_label() takes exactly 1 positional argument (4 given)'], 'median_reward': ['EXC TypeError:median_reward() takes exactly 1 positional argument (9 given)', 'EXC TypeError:median_reward() takes exactly 1 positional argument (9 given)'], 'cumulative_max': ['EXC TypeError:cumulative_max() takes exactly 1 positional argument (6 given)', 'EXC TypeError:cumulative_max() takes exactly 1 positional argument (5 given)'], 'first_nonzero_index': ['EXC TypeError:first_nonzero_index() takes exactly 1 positional argument (6 given)', 'EXC TypeError:first_nonzero_index() takes exactly 1 positional argument (4 given)']}
total funcs tested with fails: 7
[stdout]
FAILS: {'whiten': ['max diff 0.03969073295593262', 'max diff 0.03177332878112793'], 'masked_whiten': ['max diff 0.1814074069261551', 'max diff 0.09044408798217773'], 'argmax_tokens': ['EXC TypeError:argmax_tokens() takes exactly 1 positional argument (6 given)', 'EXC TypeError:argmax_tokens() takes exactly 1 positional argument (4 given)'], 'mode_label': ['EXC TypeError:mode_label() takes exactly 1 positional argument (6 given)', 'EXC TypeError:mode_label() takes exactly 1 positional argument (4 given)'], 'median_reward': ['EXC TypeError:median_reward() takes exactly 1 positional argument (9 given)', 'EXC TypeError:median_reward() takes exactly 1 positional argument (9 given)'], 'cumulative_max': ['EXC TypeError:cumulative_max() takes exactly 1 positional argument (6 given)', 'EXC TypeError:cumulative_max() takes exactly 1 positional argument (5 given)'], 'first_nonzero_index': ['EXC TypeError:first_nonzero_index() takes exactly 1 positional argument (6 given)', 'EXC TypeError:first_nonzero_index() takes exactly 1 positional argument (4 given)']}
total funcs tested with fails: 7Re-probe whiten std convention
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
v=torch.randn(3,4,generator=g); m=(torch.rand(3,4,generator=g)>0.3).float()
tgt=R.whiten(v,m,True)
n=m.sum()
mean=(v*m).sum()/n
# candidate variances
var_b=((v-mean)**2*m).sum()/n
var_u=((v-mean)**2*m).sum()/(n-1)
for name,var in [('biased/n',var_b),('unbiased/n-1',var_u)]:
for eps in [1e-8,1e-6]:
cand=(v-mean)/(var.sqrt()+eps)
print(name,eps,'match',torch.allclose(cand,tgt,atol=1e-6),'d',(cand-tgt).abs().max().item())
print('tgt[0]:',tgt[0].tolist())
"biased/n 1e-08 match True d 0.0 biased/n 1e-06 match True d 1.1920928955078125e-06 unbiased/n-1 1e-08 match False d 0.13595354557037354 unbiased/n-1 1e-06 match False d 0.13595449924468994 tgt[0]: [1.3455535173416138, -0.08751195669174194, -1.560368537902832, 0.5857791304588318] [stdout] biased/n 1e-08 match True d 0.0 biased/n 1e-06 match True d 1.1920928955078125e-06 unbiased/n-1 1e-08 match False d 0.13595354557037354 unbiased/n-1 1e-06 match False d 0.13595449924468994 tgt[0]: [1.3455535173416138, -0.08751195669174194, -1.560368537902832, 0.5857791304588318]
Find failing whiten seed
python3 -c "
import torch, reforacle as R, objectives as O
for seed in range(20):
g=torch.Generator().manual_seed(seed)
B,T=3,4
v=torch.randn(B,T,generator=g); m=(torch.rand(B,T,generator=g)>0.3).float()
r=R.whiten(v,m,True); o=O.whiten(v,m,True)
if not torch.allclose(r,o,atol=1e-6):
print('seed',seed,'FAIL d',(r-o).abs().max().item())
n=m.sum()
mean=(v*m).sum()/n
var_b=((v-mean)**2*m).sum()/n
# maybe reforacle uses unbiased masked variance
var_u=((v-mean)**2*m).sum()/(n-1)
print(' biased match',torch.allclose((v-mean)/(var_b.sqrt()+1e-8),r,atol=1e-6))
print(' unbiased match',torch.allclose((v-mean)/(var_u.sqrt()+1e-8),r,atol=1e-6))
break
else: print('all pass')
"all pass [stdout] all pass
Corrected comprehensive test incl composites
python3 -c "
import torch, reforacle as R, objectives as O
import random
def close(a,b,atol,rtol):
a=torch.as_tensor(a); b=torch.as_tensor(b)
if a.shape!=b.shape: return False, f'shape {a.shape} vs {b.shape}'
if a.dtype==torch.bool or b.dtype==torch.bool:
return bool((a==b).all()), 'bool mismatch'
ok=torch.allclose(a.float(),b.float(),atol=atol,rtol=rtol)
return ok, '' if ok else f'max diff {(a.float()-b.float()).abs().max().item()}'
fails={}
def chk(name,args,atol=1e-6,rtol=1e-5):
if not isinstance(args,tuple): args=(args,)
try:
r=getattr(R,name)(*args); o=getattr(O,name)(*args)
except Exception as e:
fails.setdefault(name,[]).append(f'EXC {type(e).__name__}:{e}'); return
ok,msg=close(r,o,atol,rtol)
if not ok: fails.setdefault(name,[]).append((name,msg))
for seed in range(40):
g=torch.Generator().manual_seed(seed+100)
def rn(*s): return torch.randn(*s,generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
def rm(*s): return (torch.rand(*s,generator=g)>0.3).float()
B,T,V=random.Random(seed).choice([(2,3,5),(4,3,6),(3,5,4),(6,2,7),(8,4,3)])
v=rn(B,T); m=rm(B,T)
chk('whiten',(v,m,True),1e-5,1e-4); chk('whiten',(v,m,False),1e-5,1e-4)
chk('masked_whiten',(v,m,True),1e-5,1e-4); chk('masked_whiten',(v,m,False),1e-5,1e-4)
logits=rn(B,T,V); labels=ri(V,B,T)
chk('argmax_tokens',logits); chk('mode_label',labels)
chk('median_reward',rn(9)); chk('cumulative_max',rn(6)); chk('cumulative_max',rn(5,3))
chk('first_nonzero_index',rm(B,T)); chk('first_nonzero_index',rm(7))
# composites
gs=random.Random(seed).choice([2,4])
Bc=gs*random.Random(seed).choice([2,3])
lg=rn(Bc,T,V); ol=rn(Bc,T,V); rf=rn(Bc,T,V); lb=ri(V,Bc,T); cm=rm(Bc,T); rew=rn(Bc)
for est in ['k1','k2','k3']:
chk('grpo_objective',(lg,ol,rf,lb,cm,rew,gs,0.05,0.2,0.2,True,est),1e-5,1e-4)
chk('grpo_objective',(lg,ol,rf,lb,cm,rew,gs,0.1,0.15,0.25,False,est),1e-5,1e-4)
chk('rloo_objective',(lg,ol,lb,cm,rew,gs,0.2,0.2),1e-5,1e-4)
chk('reverse_kl_objective',(lg,rf,lb,cm,rn(Bc),0.1),1e-5,1e-4)
Tt=random.Random(seed).choice([4,5,6])
chk('ppo_objective',(rn(Tt),rn(Tt),rn(Tt),rn(Tt),rn(Tt),rn(1),0.99,0.95,0.2,0.2,0.2,0.5),1e-5,1e-4)
chk('ppo_objective',(rn(Tt,2),rn(Tt,2),rn(Tt,2),rn(Tt,2),rn(Tt,2),rn(2),0.9,0.9,0.3,0.15,0.1,1.0),1e-5,1e-4)
# dpo_sequence
def mk(): return rn(Bc,T,V)
chk('dpo_sequence_loss',(mk(),mk(),mk(),mk(),ri(V,Bc,T),ri(V,Bc,T),rm(Bc,T),rm(Bc,T),0.1,0.05),1e-5,1e-4)
print('FAILS:', {k:v[:2] for k,v in fails.items()} if fails else 'NONE')
"FAILS: {'whiten': [('whiten', 'max diff 0.007350623607635498'), ('whiten', 'max diff 0.05743896961212158')], 'masked_whiten': [('masked_whiten', 'max diff 0.154306098818779'), ('masked_whiten', 'max diff 0.5875788927078247')], 'mode_label': [('mode_label', 'max diff 3.0'), ('mode_label', 'max diff 5.0')]}
[stdout]
FAILS: {'whiten': [('whiten', 'max diff 0.007350623607635498'), ('whiten', 'max diff 0.05743896961212158')], 'masked_whiten': [('masked_whiten', 'max diff 0.154306098818779'), ('masked_whiten', 'max diff 0.5875788927078247')], 'mode_label': [('mode_label', 'max diff 3.0'), ('mode_label', 'max diff 5.0')]}Brute force whiten formula
python3 -c "
import torch, reforacle as R
for seed in range(30):
g=torch.Generator().manual_seed(seed+100)
B,T=6,2
v=torch.randn(B,T,generator=g); m=(torch.rand(B,T,generator=g)>0.3).float()
tgt=R.whiten(v,m,True)
n=m.sum()
mean=(v*m).sum()/n
sq=((v-mean)**2*m).sum()
cands={
'bias_n_e8':(v-mean)/((sq/n).sqrt()+1e-8),
'unbias_e8':(v-mean)/((sq/(n-1)).sqrt()+1e-8),
'bias_n_e5':(v-mean)/((sq/n).sqrt()+1e-5),
}
res={k:torch.allclose(c,tgt,atol=1e-6) for k,c in cands.items()}
if not any(res.values()):
print('seed',seed,'n',n.item(),res, 'diffs',{k:(c-tgt).abs().max().item() for k,c in cands.items()})
" | head -20(Bash completed with no output)
Which whiten formula matches per seed
python3 -c "
import torch, reforacle as R
from collections import Counter
c=Counter()
for seed in range(50):
g=torch.Generator().manual_seed(seed+100)
B,T=6,2
v=torch.randn(B,T,generator=g); m=(torch.rand(B,T,generator=g)>0.3).float()
tgt=R.whiten(v,m,True)
n=m.sum()
mean=(v*m).sum()/n
sq=((v-mean)**2*m).sum()
cands={
'bias_n':(v-mean)/((sq/n).sqrt()+1e-8),
'unbias':(v-mean)/((sq/(n-1)).sqrt()+1e-8),
}
matched=[k for k,cc in cands.items() if torch.allclose(cc,tgt,atol=1e-6)]
c[tuple(matched)]+=1
print(c)
"Counter({('bias_n',): 50})
[stdout]
Counter({('bias_n',): 50})Check per-row whiten
python3 -c "
import torch, reforacle as R, objectives as O
import random
for seed in range(40):
g=torch.Generator().manual_seed(seed+100)
def rn(*s): return torch.randn(*s,generator=g)
def rm(*s): return (torch.rand(*s,generator=g)>0.3).float()
B,T,V=random.Random(seed).choice([(2,3,5),(4,3,6),(3,5,4),(6,2,7),(8,4,3)])
v=rn(B,T); m=rm(B,T)
r=R.whiten(v,m,True); o=O.whiten(v,m,True)
if not torch.allclose(r,o,atol=1e-5):
print('seed',seed,'shape',(B,T),'d',(r-o).abs().max().item())
# test per-row whitening
n=m.sum(-1,keepdim=True); mean=(v*m).sum(-1,keepdim=True)/n
sq=((v-mean)**2*m).sum(-1,keepdim=True)
prow=(v-mean)/((sq/n).sqrt()+1e-8)
print(' per-row match', torch.allclose(prow,r,atol=1e-5))
break
"(Bash completed with no output)
Reproduce whiten fail in isolation
python3 -c "
import torch, reforacle as R, objectives as O
import random
def close(a,b,atol,rtol):
a=torch.as_tensor(a); b=torch.as_tensor(b)
ok=torch.allclose(a.float(),b.float(),atol=atol,rtol=rtol)
return ok,(a.float()-b.float()).abs().max().item()
for seed in range(40):
g=torch.Generator().manual_seed(seed+100)
def rn(*s): return torch.randn(*s,generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
def rm(*s): return (torch.rand(*s,generator=g)>0.3).float()
B,T,V=random.Random(seed).choice([(2,3,5),(4,3,6),(3,5,4),(6,2,7),(8,4,3)])
v=rn(B,T); m=rm(B,T)
for sm in [True,False]:
ok,d=close(R.whiten(v,m,sm),O.whiten(v,m,sm),1e-5,1e-4)
if not ok: print('whiten seed',seed,(B,T),'sm',sm,'d',d)
ok,d=close(R.masked_whiten(v,m,True),O.masked_whiten(v,m,True),1e-5,1e-4)
if not ok: print('masked_whiten seed',seed,(B,T),'d',d)
"whiten seed 0 (6, 2) sm False d 0.007350623607635498 whiten seed 1 (4, 3) sm False d 0.05743896961212158 whiten seed 2 (2, 3) sm False d 0.5801899433135986 whiten seed 3 (4, 3) sm False d 0.05369746685028076 whiten seed 4 (4, 3) sm False d 0.2574244737625122 whiten seed 5 (8, 4) sm False d 0.013456642627716064 whiten seed 6 (8, 4) sm False d 0.010870099067687988 whiten seed 7 (3, 5) sm False d 0.019872963428497314 whiten seed 8 (4, 3) sm False d 0.07706671953201294 whiten seed 9 (6, 2) sm False d 0.08643913269042969 whiten seed 10 (8, 4) sm False d 0.028693437576293945 whiten seed 11 (6, 2) sm False d 0.01796942949295044 whiten seed 12 (6, 2) sm False d 0.0006452798843383789 whiten seed 13 (3, 5) sm False d 0.3139948844909668 whiten seed 14 (2, 3) sm False d 0.2315637171268463 whiten seed 15 (4, 3) sm False d 0.015336275100708008 whiten seed 16 (3, 5) sm False d 0.04236173629760742 whiten seed 17 (8, 4) sm False d 0.06450235843658447 whiten seed 18 (4, 3) sm False d 0.04302293062210083 whiten seed 20 (4, 3) sm False d 0.14050915837287903 whiten seed 21 (4, 3) sm False d 0.21168136596679688 whiten seed 22 (4, 3) sm False d 0.09615707397460938 whiten seed 23 (3, 5) sm False d 0.033590078353881836 whiten seed 24 (6, 2) sm False d 0.1330336332321167 whiten seed 25 (6, 2) sm False d 0.04641556739807129 whiten seed 26 (4, 3) sm False d 0.05626857280731201 whiten seed 27 (6, 2) sm False d 0.0017248988151550293 whiten seed 28 (2, 3) sm False d 0.07809117436408997 whiten seed 29 (8, 4) sm False d 0.000544428825378418 whiten seed 30 (8, 4) sm False d 0.043590664863586426 whiten seed 31 (2, 3) sm False d 2.207266092300415 whiten seed 32 (2, 3) sm False d 0.005549073219299316 whiten seed 33 (8, 4) sm False d 0.023572564125061035 whiten seed 34 (8, 4) sm False d 0.05337238311767578 whiten seed 35 (8, 4) sm False d 0.01291799545288086 whiten seed 36 (3, 5) sm False d 0.07826721668243408 whiten seed 37 (8, 4) sm False d 0.01590251922607422 whiten seed 38 (6, 2) sm False d 0.061852216720581055 whiten seed 39 (4, 3) sm False d 0.016144931316375732 [stdout] whiten seed 0 (6, 2) sm False d 0.007350623607635498 whiten seed 1 (4, 3) sm False d 0.05743896961212158 whiten seed 2 (2, 3) sm False d 0.5801899433135986 whiten seed 3 (4, 3) sm False d 0.05369746685028076 whiten seed 4 (4, 3) sm False d 0.2574244737625122 whiten seed 5 (8, 4) sm False d 0.013456642627716064 whiten seed 6 (8, 4) sm False d 0.010870099067687988 whiten seed 7 (3, 5) sm False d 0.019872963428497314 whiten seed 8 (4, 3) sm False d 0.07706671953201294 whiten seed 9 (6, 2) sm False d 0.08643913269042969 whiten seed 10 (8, 4) sm False d 0.028693437576293945 whiten seed 11 (6, 2) sm False d 0.01796942949295044 whiten seed 12 (6, 2) sm False d 0.0006452798843383789 whiten seed 13 (3, 5) sm False d 0.3139948844909668 whiten seed 14 (2, 3) sm False d 0.2315637171268463 whiten seed 15 (4, 3) sm False d 0.015336275100708008 whiten seed 16 (3, 5) sm False d 0.04236173629760742 whiten seed 17 (8, 4) sm False d 0.06450235843658447 whiten seed 18 (4, 3) sm False d 0.04302293062210083 whiten seed 20 (4, 3) sm False d 0.14050915837287903 whiten seed 21 (4, 3) sm False d 0.21168136596679688 whiten seed 22 (4, 3) sm False d 0.09615707397460938 whiten seed 23 (3, 5) sm False d 0.033590078353881836 whiten seed 24 (6, 2) sm False d 0.1330336332321167 whiten seed 25 (6, 2) sm False d 0.04641556739807129 whiten seed 26 (4, 3) sm False d 0.05626857280731201 whiten seed 27 (6, 2) sm False d 0.0017248988151550293 whiten seed 28 (2, 3) sm False d 0.07809117436408997 whiten seed 29 (8, 4) sm False d 0.000544428825378418 whiten seed 30 (8, 4) sm False d 0.043590664863586426 whiten seed 31 (2, 3) sm False d 2.207266092300415 whiten seed 32 (2, 3) sm False d 0.005549073219299316 whiten seed 33 (8, 4) sm False d 0.023572564125061035 whiten seed 34 (8, 4) sm False d 0.05337238311767578 whiten seed 35 (8, 4) sm False d 0.01291799545288086 whiten seed 36 (3, 5) sm False d 0.07826721668243408 whiten seed 37 (8, 4) sm False d 0.01590251922607422 whiten seed 38 (6, 2) sm False d 0.061852216720581055 whiten seed 39 (4, 3) sm False d 0.016144931316375732
Verify TRL add-mean-back convention
python3 -c "
import torch, reforacle as R
for seed in [0,1,2,31]:
g=torch.Generator().manual_seed(seed+100)
B,T=6,2
v=torch.randn(B,T,generator=g); m=(torch.rand(B,T,generator=g)>0.3).float()
tgt=R.whiten(v,m,False)
n=m.sum(); mean=(v*m).sum()/n; var=((v-mean)**2*m).sum()/n
add_mean=(v-mean)*torch.rsqrt(var+1e-8)+mean
add_mean2=(v-mean)/(var.sqrt()+1e-8)+mean
print(seed,'add_mean rsqrt',torch.allclose(add_mean,tgt,atol=1e-6),'sqrt+eps',torch.allclose(add_mean2,tgt,atol=1e-6))
# masked_whiten sm=False
print('--- masked_whiten sm=False ---')
for seed in [0,1,2]:
g=torch.Generator().manual_seed(seed+100)
B,T=6,2
v=torch.randn(B,T,generator=g); m=(torch.rand(B,T,generator=g)>0.3).float()
tgt=R.masked_whiten(v,m,False)
n=m.sum(); mean=(v*m).sum()/n; var=((v-mean)**2*m).sum()/n
wm=((v-mean)*torch.rsqrt(var+1e-8)+mean)
print(seed,'no mask mult',torch.allclose(wm,tgt,atol=1e-6),'w/ mask mult',torch.allclose(wm*m,tgt,atol=1e-6))
"0 add_mean rsqrt True sqrt+eps True 1 add_mean rsqrt True sqrt+eps True 2 add_mean rsqrt True sqrt+eps True 31 add_mean rsqrt True sqrt+eps True --- masked_whiten sm=False --- 0 no mask mult False w/ mask mult False 1 no mask mult False w/ mask mult False 2 no mask mult False w/ mask mult False [stdout] 0 add_mean rsqrt True sqrt+eps True 1 add_mean rsqrt True sqrt+eps True 2 add_mean rsqrt True sqrt+eps True 31 add_mean rsqrt True sqrt+eps True --- masked_whiten sm=False --- 0 no mask mult False w/ mask mult False 1 no mask mult False w/ mask mult False 2 no mask mult False w/ mask mult False
Brute masked_whiten sm=False
python3 -c "
import torch, reforacle as R
seed=0
g=torch.Generator().manual_seed(seed+100)
B,T=6,2
v=torch.randn(B,T,generator=g); m=(torch.rand(B,T,generator=g)>0.3).float()
tgt=R.masked_whiten(v,m,False)
print('tgt:',tgt.reshape(-1).tolist())
print('mask:',m.reshape(-1).tolist())
n=m.sum(); mean=(v*m).sum()/n; var=((v-mean)**2*m).sum()/n
base=(v-mean)*torch.rsqrt(var+1e-8)
cands={
'base+mean, no mask': base+mean,
'(base+mean)*mask': (base+mean)*m,
'base*mask+mean': base*m+mean,
'base*mask+mean*mask': (base+mean)*m,
'base+mean then *mask': (base+mean)*m,
}
for k,c in cands.items():
print(k, torch.allclose(c,tgt,atol=1e-5), (c-tgt).abs().max().item())
print('base:',base.reshape(-1).tolist())
print('base+mean:',(base+mean).reshape(-1).tolist())
"tgt: [0.3361250162124634, -0.2796515226364136, -0.38240861892700195, 0.22401471436023712, -0.154306098818779, -2.210582971572876, -0.30943578481674194, -0.154306098818779, 1.657562494277954, -0.2700720429420471, -0.154306098818779, -0.154306098818779] mask: [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0] base+mean, no mask False 1.1704715490341187 (base+mean)*mask False 0.154306098818779 base*mask+mean True 2.9802322387695312e-08 base*mask+mean*mask False 0.154306098818779 base+mean then *mask False 0.154306098818779 base: [0.49043112993240356, -0.1253453940153122, -0.22810250520706177, 0.3783207833766937, -1.1704715490341187, -2.056276798248291, -0.15512967109680176, -0.6778357028961182, 1.8118685483932495, -0.11576594412326813, -0.7820638418197632, 0.6031266450881958] base+mean: [0.3361250162124634, -0.2796514928340912, -0.38240861892700195, 0.22401468455791473, -1.324777603149414, -2.210582971572876, -0.30943578481674194, -0.8321418166160583, 1.657562494277954, -0.2700720429420471, -0.9363699555397034, 0.4488205313682556] [stdout] tgt: [0.3361250162124634, -0.2796515226364136, -0.38240861892700195, 0.22401471436023712, -0.154306098818779, -2.210582971572876, -0.30943578481674194, -0.154306098818779, 1.657562494277954, -0.2700720429420471, -0.154306098818779, -0.154306098818779] mask: [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0] base+mean, no mask False 1.1704715490341187 (base+mean)*mask False 0.154306098818779 base*mask+mean True 2.9802322387695312e-08 base*mask+mean*mask False 0.154306098818779 base+mean then *mask False 0.154306098818779 base: [0.49043112993240356, -0.1253453940153122, -0.22810250520706177, 0.3783207833766937, -1.1704715490341187, -2.056276798248291, -0.15512967109680176, -0.6778357028961182, 1.8118685483932495, -0.11576594412326813, -0.7820638418197632, 0.6031266450881958] base+mean: [0.3361250162124634, -0.2796514928340912, -0.38240861892700195, 0.22401468455791473, -1.324777603149414, -2.210582971572876, -0.30943578481674194, -0.8321418166160583, 1.657562494277954, -0.2700720429420471, -0.9363699555397034, 0.4488205313682556]
Probe mode_label tie-breaking
python3 -c " import torch, reforacle as R # mode with ties for lab in [torch.tensor([0,1,2,3]), torch.tensor([2,2,1,1]), torch.tensor([3,3,0,0,1]), torch.tensor([1,1,2,2,0,0])]: print(lab.tolist(),'reforacle->',R.mode_label(lab).item(),' torch.mode->',lab.mode(-1).values.item(),' torch.bincount.argmax->',torch.bincount(lab).argmax().item()) "
[0, 1, 2, 3] reforacle-> 3 torch.mode-> 0 torch.bincount.argmax-> 0 [2, 2, 1, 1] reforacle-> 2 torch.mode-> 1 torch.bincount.argmax-> 1 [3, 3, 0, 0, 1] reforacle-> 3 torch.mode-> 0 torch.bincount.argmax-> 0 [1, 1, 2, 2, 0, 0] reforacle-> 2 torch.mode-> 0 torch.bincount.argmax-> 0 [stdout] [0, 1, 2, 3] reforacle-> 3 torch.mode-> 0 torch.bincount.argmax-> 0 [2, 2, 1, 1] reforacle-> 2 torch.mode-> 1 torch.bincount.argmax-> 1 [3, 3, 0, 0, 1] reforacle-> 3 torch.mode-> 0 torch.bincount.argmax-> 0 [1, 1, 2, 2, 0, 0] reforacle-> 2 torch.mode-> 0 torch.bincount.argmax-> 0
/app/objectives.py
def _whiten(values, mask, shift_mean):
if mask is None:
mean = values.mean()
var = ((values - mean) ** 2).mean()
else:
den = mask.sum().clamp(min=1e-8)
mean = (values * mask).sum() / den
var = ((values - mean) ** 2 * mask).sum() / den
std = var.sqrt()
out = values / (std + 1e-8) if not shift_mean else (values - mean) / (std + 1e-8)
return out
def whiten(values, mask, shift_mean):
return _whiten(values, mask, shift_mean)def _whiten(values, mask, shift_mean, apply_mask=False):
if mask is None:
mean = values.mean()
var = ((values - mean) ** 2).mean()
else:
den = mask.sum().clamp(min=1e-8)
mean = (values * mask).sum() / den
var = ((values - mean) ** 2 * mask).sum() / den
out = (values - mean) * torch.rsqrt(var + 1e-8)
if apply_mask and mask is not None:
out = out * mask
if not shift_mean:
out = out + mean
return out
def whiten(values, mask, shift_mean):
return _whiten(values, mask, shift_mean, apply_mask=False)/app/objectives.py
def masked_whiten(values, mask, shift_mean):
return _whiten(values, mask, shift_mean) * maskdef masked_whiten(values, mask, shift_mean):
return _whiten(values, mask, shift_mean, apply_mask=True)/app/objectives.py
def mode_label(labels):
return labels.mode(-1).valuesdef mode_label(labels):
counts = (labels.unsqueeze(-1) == labels.unsqueeze(-2)).sum(-1)
lo = labels.min()
span = (labels.max() - lo + 1)
score = counts * span + (labels - lo)
idx = score.argmax(-1)
return labels.gather(-1, idx.unsqueeze(-1)).squeeze(-1)Full comprehensive re-test
python3 -c "
import torch, reforacle as R, objectives as O
import random
import importlib; importlib.reload(O)
def close(a,b,atol,rtol):
a=torch.as_tensor(a); b=torch.as_tensor(b)
if a.shape!=b.shape: return False, f'shape {a.shape} vs {b.shape}'
if a.dtype==torch.bool or b.dtype==torch.bool or a.dtype==torch.int64:
return bool((a==b).all()), 'int/bool mismatch'
return torch.allclose(a.float(),b.float(),atol=atol,rtol=rtol), f'max diff {(a.float()-b.float()).abs().max().item()}'
fails={}
def chk(name,args,atol=1e-6,rtol=1e-5):
if not isinstance(args,tuple): args=(args,)
try:
r=getattr(R,name)(*args); o=getattr(O,name)(*args)
except Exception as e:
fails.setdefault(name,[]).append(f'EXC {type(e).__name__}:{e}'); return
ok,msg=close(r,o,atol,rtol)
if not ok: fails.setdefault(name,[]).append(msg)
A=1e-5; RT=1e-4
for seed in range(60):
g=torch.Generator().manual_seed(seed+7)
def rn(*s): return torch.randn(*s,generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
def rm(*s): return (torch.rand(*s,generator=g)>0.3).float()
B,T,V=random.Random(seed).choice([(2,3,5),(4,3,6),(3,5,4),(6,2,7),(8,4,3),(5,6,8)])
v=rn(B,T); m=rm(B,T)
chk('masked_mean',(v,m)); chk('masked_mean',(v,m,0)); chk('masked_mean',(v,m,1))
chk('masked_sum',(v,m)); chk('masked_sum',(v,m,1)); chk('masked_sum',(v,m,0))
chk('logsumexp',(rn(B,T),1)); chk('log_softmax',(rn(B,T),1))
logits=rn(B,T,V); labels=ri(V,B,T)
chk('selective_logprobs',(logits,labels,m)); chk('token_logprobs',(logits,labels))
chk('sequence_logprob',(logits,labels,m,True)); chk('sequence_logprob',(logits,labels,m,False))
chk('entropy',(logits,m)); chk('normalized_entropy',(logits,m))
pc,pr,rc,rr=rn(B),rn(B),rn(B),rn(B)
chk('dpo_loss',(pc,pr,rc,rr,0.1,0.0)); chk('dpo_loss',(pc,pr,rc,rr,0.5,0.1))
chk('ipo_loss',(pc,pr,rc,rr,0.3)); chk('bradley_terry_logit',(rc,rr,0.2))
gs=random.Random(seed).choice([2,3,4]); rew=rn(gs*3)
chk('grpo_advantages',(rew,gs,True)); chk('grpo_advantages',(rew,gs,False))
chk('rloo_advantages',(rew,gs)); chk('group_mean_baseline',(rew,gs))
Tt=random.Random(seed).choice([4,5,6,7])
chk('gae',(rn(Tt),rn(Tt),rn(1),0.99,0.95),A,RT); chk('gae',(rn(Tt,2),rn(Tt,2),rn(2),0.9,0.9),A,RT)
chk('lambda_returns',(rn(Tt),rn(Tt),rn(1),0.95,0.9),A,RT)
chk('discounted_returns',(rn(Tt,2),0.9),A,RT); chk('discounted_returns',(rn(Tt),0.95),A,RT)
lp,ref=rn(B,T),rn(B,T)
for est in ['k1','k2','k3']: chk('kl_penalty',(lp,ref,est))
chk('reverse_kl',(lp,ref)); chk('symmetric_kl',(lp,ref))
chk('clipped_pg_loss',(lp,rn(B,T),rn(B,T),m,0.2,0.3),A,RT)
chk('value_loss',(rn(B,T),rn(B,T),rn(B,T),0.2),A,RT)
chk('whiten',(v,m,True),A,RT); chk('whiten',(v,m,False),A,RT)
chk('masked_whiten',(v,m,True),A,RT); chk('masked_whiten',(v,m,False),A,RT)
chk('importance_ratio',(lp,rn(B,T),0.2)); chk('importance_ratio',(lp,rn(B,T),None))
chk('normalize',(rn(B,T),1e-5))
chk('smoothed_nll',(logits,labels,0.1)); chk('smoothed_nll',(logits,labels,0.2))
chk('cross_entropy',(logits,labels,-100))
lab2=labels.clone(); lab2[0,0]=-100; chk('cross_entropy',(logits,lab2,-100))
chk('logprob_at_temperature',(logits,labels,2.0)); chk('logprob_at_temperature',(logits,labels,0.7))
chk('clip_fraction',(lp,rn(B,T),0.2))
chk('advantage_mean_std',(rn(B,T),m),A,RT)
chk('argmax_tokens',(logits,)); chk('mode_label',(labels,)); chk('mode_label',(ri(V,9),))
chk('median_reward',(rn(9),)); chk('masked_max',(v,m))
chk('quantile_lower',(rn(11),0.3)); chk('quantile_lower',(rn(11),0.75))
chk('huber_value_loss',(rn(B,T),rn(B,T),1.0),A,RT); chk('huber_value_loss',(rn(B,T),rn(B,T),0.5),A,RT)
chk('cumulative_max',(rn(6),)); chk('cumulative_max',(rn(5,3),))
chk('pad_mask_from_lengths',(ri(T+1,B),T))
chk('first_nonzero_index',(rm(B,T),)); chk('first_nonzero_index',(rm(7),))
probs=torch.softmax(rn(B,V),-1); chk('top_p_mask',(probs,0.7),A,RT); chk('top_p_mask',(probs,0.9),A,RT)
chk('top_k_mask',(rn(B,V),2)); chk('top_k_mask',(rn(B,V),3))
bnd,_=torch.sort(rn(3)); chk('bucketize_reward',(rn(8),bnd))
# composites
Bc=gs*random.Random(seed).choice([2,3])
lg=rn(Bc,T,V); ol=rn(Bc,T,V); rf=rn(Bc,T,V); lb=ri(V,Bc,T); cm=rm(Bc,T); rw=rn(Bc)
for est in ['k1','k2','k3']:
chk('grpo_objective',(lg,ol,rf,lb,cm,rw,gs,0.05,0.2,0.2,True,est),A,RT)
chk('grpo_objective',(lg,ol,rf,lb,cm,rw,gs,0.1,0.15,0.25,False,est),A,RT)
chk('rloo_objective',(lg,ol,lb,cm,rw,gs,0.2,0.2),A,RT)
chk('reverse_kl_objective',(lg,rf,lb,cm,rn(Bc),0.1),A,RT)
chk('ppo_objective',(rn(Tt),rn(Tt),rn(Tt),rn(Tt),rn(Tt),rn(1),0.99,0.95,0.2,0.2,0.2,0.5),A,RT)
chk('ppo_objective',(rn(Tt,2),rn(Tt,2),rn(Tt,2),rn(Tt,2),rn(Tt,2),rn(2),0.9,0.9,0.3,0.15,0.1,1.0),A,RT)
def mk(): return rn(Bc,T,V)
chk('dpo_sequence_loss',(mk(),mk(),mk(),mk(),ri(V,Bc,T),ri(V,Bc,T),rm(Bc,T),rm(Bc,T),0.1,0.05),A,RT)
print('FAILS:', {k:v[:2] for k,v in fails.items()} if fails else 'ALL PASS')
"FAILS: ALL PASS [stdout] FAILS: ALL PASS
Edge case stress tests
python3 -c "
import torch, reforacle as R, objectives as O
# top_k ties
lg=torch.tensor([[1.0,3.0,3.0,2.0,5.0]])
print('topk tie k=2:', R.top_k_mask(lg,2).tolist(), O.top_k_mask(lg,2).tolist())
print('topk tie k=3:', R.top_k_mask(lg,3).tolist(), O.top_k_mask(lg,3).tolist())
# top_p full and tiny
p=torch.softmax(torch.randn(2,6),-1)
for pp in [0.0,0.1,0.5,1.0,0.99]:
r=R.top_p_mask(p,pp); o=O.top_p_mask(p,pp)
print('top_p',pp,torch.allclose(r,o,atol=1e-5), (r-o).abs().max().item())
# mode_label big
lab=torch.randint(0,10,(5,20))
print('mode big:', torch.equal(R.mode_label(lab),O.mode_label(lab)))
# bucketize exact boundary
b=torch.tensor([0.0,1.0,2.0]); rw=torch.tensor([0.0,1.0,2.0,-0.0,0.9999])
print('bucketize exact:', R.bucketize_reward(rw,b).tolist(), O.bucketize_reward(rw,b).tolist())
# masked_max empty dtype
print('masked_max empty:', R.masked_max(torch.randn(4),torch.zeros(4)), O.masked_max(torch.randn(4),torch.zeros(4)))
# quantile lower various
x=torch.randn(13)
for q in [0.0,0.25,0.5,0.9,1.0]:
print('ql',q, torch.allclose(R.quantile_lower(x,q),O.quantile_lower(x,q)))
"topk tie k=2: [[False, True, True, False, True]] [[False, True, True, False, True]] topk tie k=3: [[False, True, True, False, True]] [[False, True, True, False, True]] top_p 0.0 True 0.0 top_p 0.1 True 0.0 top_p 0.5 True 0.0 top_p 1.0 True 0.0 top_p 0.99 True 0.0 mode big: True bucketize exact: [1, 2, 3, 1, 1] [1, 2, 3, 1, 1] masked_max empty: tensor(0.) tensor(0.) ql 0.0 True ql 0.25 True ql 0.5 True ql 0.9 True ql 1.0 True [stdout] topk tie k=2: [[False, True, True, False, True]] [[False, True, True, False, True]] topk tie k=3: [[False, True, True, False, True]] [[False, True, True, False, True]] top_p 0.0 True 0.0 top_p 0.1 True 0.0 top_p 0.5 True 0.0 top_p 1.0 True 0.0 top_p 0.99 True 0.0 mode big: True bucketize exact: [1, 2, 3, 1, 1] [1, 2, 3, 1, 1] masked_max empty: tensor(0.) tensor(0.) ql 0.0 True ql 0.25 True ql 0.5 True ql 0.9 True ql 1.0 True
Float64 and 3D robustness
python3 -c "
import torch, reforacle as R, objectives as O
import random
torch.set_default_dtype(torch.float64)
fails=[]
def chk(name,args,A=1e-5,RT=1e-4):
if not isinstance(args,tuple): args=(args,)
r=getattr(R,name)(*args); o=getattr(O,name)(*args)
r=torch.as_tensor(r); o=torch.as_tensor(o)
if r.shape!=o.shape: fails.append((name,'shape',r.shape,o.shape)); return
if r.dtype in (torch.bool,torch.int64):
if not (r==o).all(): fails.append((name,'intbool')); return
elif not torch.allclose(r.double(),o.double(),atol=A,rtol=RT):
fails.append((name,(r-o).abs().max().item()))
for seed in range(30):
g=torch.Generator().manual_seed(seed*3+1)
def rn(*s): return torch.randn(*s,generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
def rm(*s): return (torch.rand(*s,generator=g)>0.4).double()
B,T,V=random.Random(seed).choice([(3,4,5),(6,3,7),(4,5,6)])
v=rn(B,T); m=rm(B,T); logits=rn(B,T,V); labels=ri(V,B,T)
for a in [None,0,1]: chk('masked_mean',(v,m,a) if a is not None else (v,m))
chk('whiten',(v,m,False)); chk('masked_whiten',(v,m,False))
chk('whiten',(v,m,True)); chk('masked_whiten',(v,m,True))
chk('mode_label',(labels,))
gs=2; Bc=4
lg=rn(Bc,T,V); ol=rn(Bc,T,V); rf=rn(Bc,T,V); lb=ri(V,Bc,T); cm=rm(Bc,T); rw=rn(Bc)
chk('grpo_objective',(lg,ol,rf,lb,cm,rw,gs,0.05,0.2,0.3,True,'k3'))
Tt=5
chk('ppo_objective',(rn(Tt),rn(Tt),rn(Tt),rn(Tt),rn(Tt),rn(1),0.99,0.95,0.2,0.2,0.2,0.5))
# 3D masked_mean
v3=rn(2,3,4); m3=rm(2,3,4)
chk('masked_mean',(v3,m3,2)); chk('masked_sum',(v3,m3,1))
print('FAILS float64:', fails if fails else 'ALL PASS')
"FAILS float64: ALL PASS [stdout] FAILS float64: ALL PASS
Verify standalone and complete
grep -n "reforacle\|import" objectives.py | head; echo "---compile---"; python3 -c "
import ast
src=open('/app/objectives.py').read()
assert 'reforacle' not in src, 'references reforacle!'
tree=ast.parse(src)
fns=[n.name for n in tree.body if isinstance(n,ast.FunctionDef)]
import objectives
stubs=['masked_mean','masked_sum','logsumexp','log_softmax','selective_logprobs','token_logprobs','sequence_logprob','entropy','dpo_loss','ipo_loss','grpo_advantages','gae','kl_penalty','clipped_pg_loss','value_loss','whiten','dpo_sequence_loss','grpo_objective','ppo_objective','rloo_advantages','reverse_kl','importance_ratio','discounted_returns','normalize','top_p_mask','smoothed_nll','bradley_terry_logit','rloo_objective','reverse_kl_objective','cross_entropy','top_k_mask','group_mean_baseline','lambda_returns','symmetric_kl','huber_value_loss','normalized_entropy','clip_fraction','masked_whiten','logprob_at_temperature','advantage_mean_std','argmax_tokens','mode_label','median_reward','masked_max','quantile_lower','pad_mask_from_lengths','first_nonzero_index','cumulative_max','bucketize_reward']
missing=[s for s in stubs if not callable(getattr(objectives,s,None))]
print('num stubs:',len(stubs),'missing:',missing)
print('all 49 implemented and standalone OK')
"1:from __future__ import annotations 3:import torch 4:import torch.nn.functional as F 186:def importance_ratio(logp, old_logp, clip): ---compile--- num stubs: 49 missing: [] all 49 implemented and standalone OK [stdout] 1:from __future__ import annotations 3:import torch 4:import torch.nn.functional as F 186:def importance_ratio(logp, old_logp, clip): ---compile--- num stubs: 49 missing: [] all 49 implemented and standalone OK
No file changes, the agent left the workspace untouched.
=== SyncVals verifier (post-training objectives, sealed oracle) === FUNCTIONS PASSED: 47/49 (partial_score=0.9592) REWARD = 0 (1 iff ALL 49 functions match the sealed reference exactly) The agent probes a sealed compiled oracle (reforacle) to recover 49 post-training objective conventions (DPO/IPO/GRPO/PPO/GAE/KL/RLOO) and reimplements them in torch.
Reproduce this trial: git checkout 2f94510 && PYTHONPATH=src python3 scripts/build_site.py , then open trial/trial_7ec11d61d92641d9. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.
Trial trial_7ec11d61d92641d9 · verifier authoritative; classifier explanatory.