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:59Run probe.py
python3 probe.py 2>&1 | head -30
inputs: values = [1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548, 0.40334683656692505, 0.8380263447761536] 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. [stdout] inputs: values = [1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548, 0.40334683656692505, 0.8380263447761536] 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. ['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] ['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']
List reforacle functions
python3 -c "import reforacle; print([x for x in dir(reforacle) if not x.startswith('_')])"inputs: values = [1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548, 0.40334683656692505, 0.8380263447761536] 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. [stdout] inputs: values = [1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548, 0.40334683656692505, 0.8380263447761536] 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. ['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] ['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']
import torch, reforacle as R
g = torch.Generator().manual_seed(1)
def rn(*s): return torch.randn(*s, generator=g)
def ru(*s): return torch.rand(*s, generator=g)
def show(name, out):
if isinstance(out, torch.Tensor):
print(f"{name}: shape={tuple(out.shape)} dtype={out.dtype}")
print(" ", out.flatten()[:12].tolist())
elif isinstance(out, (tuple, list)):
print(f"{name}: tuple/list len {len(out)}")
for i,o in enumerate(out):
show(f" [{i}]", o)
else:
print(f"{name}: {out}")
import sys
sel = sys.argv[1] if len(sys.argv)>1 else "all"
# ---- masked ops
v = rn(4,5); m = (ru(4,5)>0.4).float()
if sel in ("all","mask"):
show("masked_mean axis=None", R.masked_mean(v,m))
show("masked_mean axis=1", R.masked_mean(v,m,1))
show("masked_sum axis=None", R.masked_sum(v,m))
show("masked_sum axis=1", R.masked_sum(v,m,1))
show("masked_max", R.masked_max(v,m))
if sel in ("all","lse"):
x = rn(3,4)
show("logsumexp axis=1", R.logsumexp(x,1))
show("log_softmax axis=1", R.log_softmax(x,1))
import torch, reforacle as R
import torch.nn.functional as F
g = torch.Generator().manual_seed(1)
def rn(*s): return torch.randn(*s, generator=g)
def ru(*s): return torch.rand(*s, generator=g)
def show(name, out):
if isinstance(out, torch.Tensor):
print(f"{name}: shape={tuple(out.shape)} dtype={out.dtype}")
fl = out.flatten()
print(" ", [round(x,5) for x in fl[:10].tolist()])
elif isinstance(out, (tuple, list)):
print(f"{name}: tuple len {len(out)}")
for i,o in enumerate(out): show(f" [{i}]", o)
else:
print(f"{name}: {out!r}")
def try_(name, fn):
try: show(name, fn())
except Exception as e: print(f"{name}: ERR {type(e).__name__}: {e}")
import sys
S = sys.argv[1:] if len(sys.argv)>1 else ["all"]
def go(tag): return "all" in S or tag in S
if go("mask"):
v = rn(3,4); m = (ru(3,4)>0.4).float()
print("v=",v.tolist()); print("m=",m.tolist())
try_("masked_mean None", lambda: R.masked_mean(v,m))
try_("masked_mean 1", lambda: R.masked_mean(v,m,1))
try_("masked_sum None", lambda: R.masked_sum(v,m))
try_("masked_sum 1", lambda: R.masked_sum(v,m,1))
try_("masked_max", lambda: R.masked_max(v,m))
try_("whiten shiftT", lambda: R.whiten(v.flatten(), None if False else torch.ones(12), True))
try_("masked_whiten shiftT", lambda: R.masked_whiten(v,m,True))
try_("masked_whiten shiftF", lambda: R.masked_whiten(v,m,False))
if go("lp"):
logits = rn(2,3,5); labels = torch.randint(0,5,(2,3),generator=g); m=(ru(2,3)>0.3).float()
try_("token_logprobs", lambda: R.token_logprobs(logits,labels))
try_("selective_logprobs", lambda: R.selective_logprobs(logits,labels,m))
try_("sequence_logprob F", lambda: R.sequence_logprob(logits,labels,m,False))
try_("sequence_logprob T", lambda: R.sequence_logprob(logits,labels,m,True))
try_("entropy", lambda: R.entropy(logits,m))
try_("normalized_entropy", lambda: R.normalized_entropy(logits,m))
try_("cross_entropy ii-1", lambda: R.cross_entropy(logits,labels,-100))
try_("smoothed_nll", lambda: R.smoothed_nll(logits,labels,0.1))
try_("logprob_at_temperature", lambda: R.logprob_at_temperature(logits,labels,0.7))
try_("argmax_tokens", lambda: R.argmax_tokens(logits))
if go("pref"):
pc,pr,rc,rr = rn(4),rn(4),rn(4),rn(4)
try_("dpo ls0", lambda: R.dpo_loss(pc,pr,rc,rr,0.1,0.0))
try_("dpo ls.1", lambda: R.dpo_loss(pc,pr,rc,rr,0.1,0.1))
try_("ipo", lambda: R.ipo_loss(pc,pr,rc,rr,0.1))
try_("bt_logit", lambda: R.bradley_terry_logit(rn(4),rn(4),0.1))
if go("adv"):
rew = rn(6)
try_("grpo_adv stdT", lambda: R.grpo_advantages(rew,3,True))
try_("grpo_adv stdF", lambda: R.grpo_advantages(rew,3,False))
try_("rloo_adv", lambda: R.rloo_advantages(rew,3))
try_("group_mean_baseline", lambda: R.group_mean_baseline(rew,3))
try_("advantage_mean_std", lambda: R.advantage_mean_std(rew,(ru(6)>0.3).float()))
r2=rn(4,5)
try_("gae", lambda: R.gae(r2, rn(4,5), rn(4), 0.99,0.95))
try_("lambda_returns", lambda: R.lambda_returns(r2, rn(4,5), rn(4),0.99,0.95))
try_("discounted_returns", lambda: R.discounted_returns(r2,0.99))
if go("kl"):
lp,rlp = rn(3,4), rn(3,4)
try_("kl k1", lambda: R.kl_penalty(lp,rlp,"k1"))
try_("kl k2", lambda: R.kl_penalty(lp,rlp,"k2"))
try_("kl k3", lambda: R.kl_penalty(lp,rlp,"k3"))
try_("reverse_kl", lambda: R.reverse_kl(lp,rlp))
try_("symmetric_kl", lambda: R.symmetric_kl(lp,rlp))
if go("pg"):
lp,olp = rn(3,4), rn(3,4); adv=rn(3,4); m=(ru(3,4)>0.3).float()
try_("clipped_pg", lambda: R.clipped_pg_loss(lp,olp,adv,m,0.2,0.2))
try_("importance_ratio noclip", lambda: R.importance_ratio(lp,olp,None))
try_("importance_ratio clip", lambda: R.importance_ratio(lp,olp,0.2))
try_("clip_fraction", lambda: R.clip_fraction(lp,olp,0.2))
val,ov,ret=rn(3,4),rn(3,4),rn(3,4)
try_("value_loss clip", lambda: R.value_loss(val,ov,ret,0.2))
try_("value_loss noclip", lambda: R.value_loss(val,ov,ret,None))
try_("huber_value_loss", lambda: R.huber_value_loss(val,ret,1.0))
if go("misc"):
try_("normalize", lambda: R.normalize(rn(5),1e-8))
p=F.softmax(rn(2,5),-1)
try_("top_p_mask", lambda: R.top_p_mask(p,0.8))
try_("top_k_mask", lambda: R.top_k_mask(rn(2,5),2))
try_("mode_label", lambda: R.mode_label(torch.tensor([1,2,2,3,3,3])))
try_("median_reward", lambda: R.median_reward(rn(5)))
try_("quantile_lower", lambda: R.quantile_lower(rn(10),0.25))
try_("pad_mask", lambda: R.pad_mask_from_lengths(torch.tensor([2,3,1]),4))
try_("first_nonzero", lambda: R.first_nonzero_index(torch.tensor([[0,0,1,0],[0,0,0,0],[1,1,0,0.]])))
try_("cumulative_max", lambda: R.cumulative_max(torch.tensor([1.,3,2,5,4])))
try_("bucketize", lambda: R.bucketize_reward(torch.tensor([0.1,0.5,0.9,1.5]),torch.tensor([0.3,0.7,1.0])))
Run full discovery
python3 discover.py all 2>&1
v= [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525], [-0.4519059658050537, -0.16613022983074188, -1.522768497467041, 0.38168391585350037], [-1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
m= [[0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [1.0, 0.0, 0.0, 0.0]]
masked_mean None: shape=() dtype=torch.float32
[-0.19784]
masked_mean 1: shape=(3,) dtype=torch.float32
[0.31664, -0.43574, -1.02761]
masked_sum None: shape=() dtype=torch.float32
[-1.3849]
masked_sum 1: shape=(3,) dtype=torch.float32
[0.94992, -1.30721, -1.02761]
masked_max: shape=() dtype=torch.float32
[0.62132]
whiten shiftT: shape=(12,) dtype=torch.float32
[1.34966, 0.74844, 0.43558, 1.28863, -0.34726, 0.08834, -1.97955, 0.92336, -1.22479, -0.51668]
masked_whiten shiftT: shape=(3, 4) dtype=torch.float32
[0.0, 0.63625, 0.35528, 1.12141, -0.0, 0.04341, -1.81378, 0.79335, -1.13592, -0.0]
masked_whiten shiftF: shape=(3, 4) dtype=torch.float32
[-0.19784, 0.43841, 0.15743, 0.92356, -0.19784, -0.15443, -2.01162, 0.59551, -1.33377, -0.19784]
token_logprobs: shape=(2, 3) dtype=torch.float32
[-1.28356, -3.76535, -1.33401, -3.14668, -2.33573, -1.75024]
selective_logprobs: shape=(2,) dtype=torch.float32
[-6.38292, -4.08597]
sequence_logprob F: shape=(2,) dtype=torch.float32
[-6.38292, -4.08597]
sequence_logprob T: shape=(2,) dtype=torch.float32
[-2.12764, -2.04299]
entropy: shape=() dtype=torch.float32
[1.23232]
normalized_entropy: shape=() dtype=torch.float32
[0.76569]
cross_entropy ii-1: shape=() dtype=torch.float32
[2.26926]
smoothed_nll: shape=() dtype=torch.float32
[2.25185]
logprob_at_temperature: shape=(2, 3) dtype=torch.float32
[-1.22024, -4.99019, -1.51029, -4.06852, -2.8164, -1.96037]
argmax_tokens: shape=(2, 3) dtype=torch.int64
[0, 3, 2, 3, 0, 3]
dpo ls0: shape=() dtype=torch.float32
[0.64596]
dpo ls.1: shape=() dtype=torch.float32
[0.65613]
ipo: shape=() dtype=torch.float32
[17.76533]
bt_logit: shape=(4,) dtype=torch.float32
[0.12859, 0.07957, -0.10346, -0.07111]
grpo_adv stdT: shape=(6,) dtype=torch.float32
[-0.13374, 1.28612, -1.15238, 1.34941, -0.30822, -1.04119]
grpo_adv stdF: shape=(6,) dtype=torch.float32
[-0.04558, 0.43827, -0.3927, 0.41313, -0.09436, -0.31877]
rloo_adv: shape=(6,) dtype=torch.float32
[-0.06836, 0.65741, -0.58905, 0.6197, -0.14155, -0.47815]
group_mean_baseline: shape=(6,) dtype=torch.float32
[-0.04558, 0.43827, -0.3927, 0.41313, -0.09436, -0.31877]
advantage_mean_std: shape=(2,) dtype=torch.float32
[-0.87102, 0.1122]
gae: ERR RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 0
lambda_returns: ERR RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 0
discounted_returns: shape=(4, 5) dtype=torch.float32
[1.71428, 0.51168, -1.52825, 0.32243, -3.38942, 1.30985, 0.73171, -0.69525, 0.74994, -2.44995]
kl k1: shape=(3, 4) dtype=torch.float32
[-0.36436, -0.22797, -0.06385, -0.82872, 0.67619, 1.17649, -0.68536, 0.78891, -1.51439, 3.07855]
kl k2: shape=(3, 4) dtype=torch.float32
[0.06638, 0.02599, 0.00204, 0.34339, 0.22862, 0.69206, 0.23486, 0.31119, 1.14669, 4.73875]
kl k3: shape=(3, 4) dtype=torch.float32
[0.07523, 0.02808, 0.00208, 0.46167, 0.18474, 0.48485, 0.29912, 0.24325, 2.03226, 2.12458]
reverse_kl: shape=(3, 4) dtype=torch.float32
[0.07523, 0.02808, 0.00208, 0.46167, 0.18474, 0.48485, 0.29912, 0.24325, 2.03226, 2.12458]
symmetric_kl: shape=(3, 4) dtype=torch.float32
[0.06712, 0.0261, 0.00204, 0.3635, 0.23746, 0.77566, 0.24419, 0.32767, 1.3833, 9.88649]
clipped_pg: shape=() dtype=torch.float32
[0.24464]
importance_ratio noclip: shape=(3, 4) dtype=torch.float32
[0.19618, 7.34988, 1.77794, 2.19257, 7.09043, 1.24907, 0.40914, 2.07743, 0.25119, 0.79662]
importance_ratio clip: shape=(3, 4) dtype=torch.float32
[0.8, 1.2, 1.2, 1.2, 1.2, 1.2, 0.8, 1.2, 0.8, 0.8]
clip_fraction: shape=() dtype=torch.float32
[1.0]
value_loss clip: shape=() dtype=torch.float32
[2.24174]
value_loss noclip: ERR TypeError: bad operand type for unary -: 'NoneType'
huber_value_loss: shape=() dtype=torch.float32
[1.36972]
normalize: shape=(5,) dtype=torch.float32
[-0.12287, -0.82484, -0.94768, 1.50005, 0.39534]
top_p_mask: shape=(2, 5) dtype=torch.float32
[0.0, 0.3141, 0.0, 0.2151, 0.4708, 0.0, 0.19083, 0.33096, 0.0, 0.47821]
top_k_mask: shape=(2, 5) dtype=torch.bool
[0, 1, 0, 0, 1, 0, 1, 1, 0, 0]
mode_label: shape=() dtype=torch.int64
[3]
median_reward: shape=() dtype=torch.float32
[-0.20429]
quantile_lower: shape=() dtype=torch.float32
[-0.30885]
pad_mask: shape=(3, 4) dtype=torch.float32
[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0]
first_nonzero: shape=(3,) dtype=torch.int64
[2, -1, 0]
cumulative_max: shape=(5,) dtype=torch.float32
[1.0, 1.0, 3.0, 3.0, 5.0]
bucketize: shape=(4,) dtype=torch.int64
[0, 1, 2, 3]
[stdout]
v= [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525], [-0.4519059658050537, -0.16613022983074188, -1.522768497467041, 0.38168391585350037], [-1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
m= [[0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [1.0, 0.0, 0.0, 0.0]]
masked_mean None: shape=() dtype=torch.float32
[-0.19784]
masked_mean 1: shape=(3,) dtype=torch.float32
[0.31664, -0.43574, -1.02761]
masked_sum None: shape=() dtype=torch.float32
[-1.3849]
masked_sum 1: shape=(3,) dtype=torch.float32
[0.94992, -1.30721, -1.02761]
masked_max: shape=() dtype=torch.float32
[0.62132]
whiten shiftT: shape=(12,) dtype=torch.float32
[1.34966, 0.74844, 0.43558, 1.28863, -0.34726, 0.08834, -1.97955, 0.92336, -1.22479, -0.51668]
masked_whiten shiftT: shape=(3, 4) dtype=torch.float32
[0.0, 0.63625, 0.35528, 1.12141, -0.0, 0.04341, -1.81378, 0.79335, -1.13592, -0.0]
masked_whiten shiftF: shape=(3, 4) dtype=torch.float32
[-0.19784, 0.43841, 0.15743, 0.92356, -0.19784, -0.15443, -2.01162, 0.59551, -1.33377, -0.19784]
token_logprobs: shape=(2, 3) dtype=torch.float32
[-1.28356, -3.76535, -1.33401, -3.14668, -2.33573, -1.75024]
selective_logprobs: shape=(2,) dtype=torch.float32
[-6.38292, -4.08597]
sequence_logprob F: shape=(2,) dtype=torch.float32
[-6.38292, -4.08597]
sequence_logprob T: shape=(2,) dtype=torch.float32
[-2.12764, -2.04299]
entropy: shape=() dtype=torch.float32
[1.23232]
normalized_entropy: shape=() dtype=torch.float32
[0.76569]
cross_entropy ii-1: shape=() dtype=torch.float32
[2.26926]
smoothed_nll: shape=() dtype=torch.float32
[2.25185]
logprob_at_temperature: shape=(2, 3) dtype=torch.float32
[-1.22024, -4.99019, -1.51029, -4.06852, -2.8164, -1.96037]
argmax_tokens: shape=(2, 3) dtype=torch.int64
[0, 3, 2, 3, 0, 3]
dpo ls0: shape=() dtype=torch.float32
[0.64596]
dpo ls.1: shape=() dtype=torch.float32
[0.65613]
ipo: shape=() dtype=torch.float32
[17.76533]
bt_logit: shape=(4,) dtype=torch.float32
[0.12859, 0.07957, -0.10346, -0.07111]
grpo_adv stdT: shape=(6,) dtype=torch.float32
[-0.13374, 1.28612, -1.15238, 1.34941, -0.30822, -1.04119]
grpo_adv stdF: shape=(6,) dtype=torch.float32
[-0.04558, 0.43827, -0.3927, 0.41313, -0.09436, -0.31877]
rloo_adv: shape=(6,) dtype=torch.float32
[-0.06836, 0.65741, -0.58905, 0.6197, -0.14155, -0.47815]
group_mean_baseline: shape=(6,) dtype=torch.float32
[-0.04558, 0.43827, -0.3927, 0.41313, -0.09436, -0.31877]
advantage_mean_std: shape=(2,) dtype=torch.float32
[-0.87102, 0.1122]
gae: ERR RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 0
lambda_returns: ERR RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 0
discounted_returns: shape=(4, 5) dtype=torch.float32
[1.71428, 0.51168, -1.52825, 0.32243, -3.38942, 1.30985, 0.73171, -0.69525, 0.74994, -2.44995]
kl k1: shape=(3, 4) dtype=torch.float32
[-0.36436, -0.22797, -0.06385, -0.82872, 0.67619, 1.17649, -0.68536, 0.78891, -1.51439, 3.07855]
kl k2: shape=(3, 4) dtype=torch.float32
[0.06638, 0.02599, 0.00204, 0.34339, 0.22862, 0.69206, 0.23486, 0.31119, 1.14669, 4.73875]
kl k3: shape=(3, 4) dtype=torch.float32
[0.07523, 0.02808, 0.00208, 0.46167, 0.18474, 0.48485, 0.29912, 0.24325, 2.03226, 2.12458]
reverse_kl: shape=(3, 4) dtype=torch.float32
[0.07523, 0.02808, 0.00208, 0.46167, 0.18474, 0.48485, 0.29912, 0.24325, 2.03226, 2.12458]
symmetric_kl: shape=(3, 4) dtype=torch.float32
[0.06712, 0.0261, 0.00204, 0.3635, 0.23746, 0.77566, 0.24419, 0.32767, 1.3833, 9.88649]
clipped_pg: shape=() dtype=torch.float32
[0.24464]
importance_ratio noclip: shape=(3, 4) dtype=torch.float32
[0.19618, 7.34988, 1.77794, 2.19257, 7.09043, 1.24907, 0.40914, 2.07743, 0.25119, 0.79662]
importance_ratio clip: shape=(3, 4) dtype=torch.float32
[0.8, 1.2, 1.2, 1.2, 1.2, 1.2, 0.8, 1.2, 0.8, 0.8]
clip_fraction: shape=() dtype=torch.float32
[1.0]
value_loss clip: shape=() dtype=torch.float32
[2.24174]
value_loss noclip: ERR TypeError: bad operand type for unary -: 'NoneType'
huber_value_loss: shape=() dtype=torch.float32
[1.36972]
normalize: shape=(5,) dtype=torch.float32
[-0.12287, -0.82484, -0.94768, 1.50005, 0.39534]
top_p_mask: shape=(2, 5) dtype=torch.float32
[0.0, 0.3141, 0.0, 0.2151, 0.4708, 0.0, 0.19083, 0.33096, 0.0, 0.47821]
top_k_mask: shape=(2, 5) dtype=torch.bool
[0, 1, 0, 0, 1, 0, 1, 1, 0, 0]
mode_label: shape=() dtype=torch.int64
[3]
median_reward: shape=() dtype=torch.float32
[-0.20429]
quantile_lower: shape=() dtype=torch.float32
[-0.30885]
pad_mask: shape=(3, 4) dtype=torch.float32
[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0]
first_nonzero: shape=(3,) dtype=torch.int64
[2, -1, 0]
cumulative_max: shape=(5,) dtype=torch.float32
[1.0, 1.0, 3.0, 3.0, 5.0]
bucketize: shape=(4,) dtype=torch.int64
[0, 1, 2, 3]import torch, reforacle as R
import torch.nn.functional as F
torch.set_printoptions(precision=6)
def eq(a,b,t=1e-5):
return torch.allclose(torch.as_tensor(a,dtype=torch.float32),torch.as_tensor(b,dtype=torch.float32),atol=t,rtol=t)
# ---- whiten vs masked_whiten; var correction
x = torch.tensor([1.,2,3,4,5,6])
m = torch.tensor([1.,1,1,0,0,1])
print("whiten shiftT:", R.whiten(x,m,True))
print("whiten shiftF:", R.whiten(x,m,False))
print("masked_whiten shiftT:", R.masked_whiten(x,m,True))
# compute candidates for whiten with mask: stats over masked
def mmean(x,m): return (x*m).sum()/m.sum()
mean=mmean(x,m)
# biased var over masked
n=m.sum()
var_b=((x-mean)**2*m).sum()/n
var_u=((x-mean)**2*m).sum()/(n-1)
print("mean",mean.item(),"var_b",var_b.item(),"var_u",var_u.item())
for eps in [1e-8,0.0]:
print("cand whiten biasedvar eps",eps, ((x-mean)*torch.rsqrt(var_b+eps)))
print("cand whiten unbiasvar eps",eps, ((x-mean)*torch.rsqrt(var_u+eps)))
# ---- entropy reduction
logits=torch.randn(2,3,5,generator=torch.Generator().manual_seed(5))
mm=torch.tensor([[1.,1,0],[1,0,0]])
p=F.softmax(logits,-1); H=-(p*F.log_softmax(logits,-1)).sum(-1)
print("entropy ref:",R.entropy(logits,mm).item())
print(" masked_mean H:", ((H*mm).sum()/mm.sum()).item())
print(" masked_sum H:", (H*mm).sum().item())
print(" plain mean H:", H.mean().item())
# ---- normalize
v=torch.tensor([1.,2,4,8])
print("normalize:",R.normalize(v,1e-8))
print(" (x-mean)/std_unb:",(v-v.mean())/(v.std()+1e-8))
print(" (x-mean)/std_bia:",(v-v.mean())/(v.std(unbiased=False)+1e-8))
print(" (x-mean)/sqrt(var_b+eps):",(v-v.mean())*torch.rsqrt(v.var(unbiased=False)+1e-8))
# ---- dpo / ipo / bt
pc,pr,rc,rr=torch.tensor([1.,0.5]),torch.tensor([0.2,0.3]),torch.tensor([0.4,0.1]),torch.tensor([-.1,.2])
beta=0.1
h=(pc-pr)-(rc-rr); lg=beta*h
print("dpo ls0 ref:",R.dpo_loss(pc,pr,rc,rr,beta,0.0).item(),"cand:",(-F.logsigmoid(lg)).mean().item())
ls=0.2
cand=(-(1-ls)*F.logsigmoid(lg)-ls*F.logsigmoid(-lg)).mean()
print("dpo ls.2 ref:",R.dpo_loss(pc,pr,rc,rr,beta,ls).item(),"cand:",cand.item())
print("ipo ref:",R.ipo_loss(pc,pr,rc,rr,beta).item(),"cand (h-1/2b)^2:",((h-1/(2*beta))**2).mean().item())
cr,rj=torch.tensor([1.,2]),torch.tensor([0.5,0.5])
print("bt ref:",R.bradley_terry_logit(cr,rj,beta))
print(" beta*(c-r):",beta*(cr-rj))
print(" logsig(beta*(c-r)):",F.logsigmoid(beta*(cr-rj)))
# ---- gae axis
T,B=4,3
rew=torch.randn(T,generator=torch.Generator().manual_seed(7))
val=torch.randn(T,generator=torch.Generator().manual_seed(8))
nv=torch.tensor(0.5)
print("gae 1D:",end=" ")
try:print(R.gae(rew,val,nv,0.99,0.95))
except Exception as e:print("ERR",e)
# manual gae 1D
def gae1d(rew,val,nv,g,l):
T=len(rew);adv=torch.zeros(T);last=0.0
for t in reversed(range(T)):
nvv=nv if t==T-1 else val[t+1]
delta=rew[t]+g*nvv-val[t]
last=delta+g*l*last;adv[t]=last
return adv
print("gae1d cand:",gae1d(rew,val,nv,0.99,0.95))
print("lambda_returns 1D:",end=" ")
try:print(R.lambda_returns(rew,val,nv,0.99,0.95))
except Exception as e:print("ERR",e)
print("lret cand adv+val:",gae1d(rew,val,nv,0.99,0.95)+val)
# discounted_returns axis
dr=torch.tensor([1.,2,3])
print("discounted 1D:",R.discounted_returns(dr,0.5))
print(" cand:",torch.tensor([1+0.5*2+0.25*3,2+0.5*3,3.]))
# ---- clip_fraction, value_loss, huber
lp=torch.tensor([0.1,0.5,-0.3,0.0]);olp=torch.tensor([0.0,0.0,0.0,0.0])
print("clip_fraction:",R.clip_fraction(lp,olp,0.2).item())
ratio=torch.exp(lp-olp)
print(" cand frac |ratio-1|>clip:",((ratio-1).abs()>0.2).float().mean().item())
val=torch.tensor([1.,2,3]);ov=torch.tensor([1.2,1.8,3.1]);ret=torch.tensor([1.5,1.5,2.5])
print("value_loss:",R.value_loss(val,ov,ret,0.2).item())
vpc=ov+torch.clamp(val-ov,-0.2,0.2)
print(" cand .5max:",(0.5*torch.max((val-ret)**2,(vpc-ret)**2)).mean().item())
print(" cand max no .5:",(torch.max((val-ret)**2,(vpc-ret)**2)).mean().item())
print("huber:",R.huber_value_loss(val,ret,1.0).item())
print(" smooth_l1:",F.smooth_l1_loss(val,ret,beta=1.0).item())
# ---- smoothed_nll, logprob_at_temperature, cross_entropy label_smoothing
lg2=torch.randn(2,3,5,generator=torch.Generator().manual_seed(11));lb=torch.randint(0,5,(2,3),generator=torch.Generator().manual_seed(12))
print("smoothed_nll:",R.smoothed_nll(lg2,lb,0.1).item())
print(" F.ce ls:",F.cross_entropy(lg2.reshape(-1,5),lb.reshape(-1),label_smoothing=0.1).item())
print("logprob_at_temp:",R.logprob_at_temperature(lg2,lb,0.7).flatten()[:4])
lpt=F.log_softmax(lg2/0.7,-1).gather(-1,lb.unsqueeze(-1)).squeeze(-1)
print(" cand:",lpt.flatten()[:4])
# ---- quantile, mode ties, median even
q=torch.tensor([1.,2,3,4,5,6,7,8,9,10])
print("quantile_lower .25:",R.quantile_lower(q,0.25).item(),"torch.quantile:",torch.quantile(q,0.25).item())
print(" quantile lower interp:",torch.quantile(q,0.25,interpolation='lower').item())
ev=torch.tensor([1.,2,3,4])
print("median_reward even:",R.median_reward(ev).item(),"torch.median:",torch.median(ev).item(),"quantile.5:",torch.quantile(ev,0.5).item())
Run second discovery
python3 d2.py 2>&1
whiten shiftT: tensor([-1.069045, -0.534522, 0.000000, 0.534522, 1.069045, 1.603567]) whiten shiftF: tensor([1.930955, 2.465477, 3.000000, 3.534523, 4.069045, 4.603567]) masked_whiten shiftT: tensor([-1.069045, -0.534522, 0.000000, 0.000000, 0.000000, 1.603567]) mean 3.0 var_b 3.5 var_u 4.666666507720947 cand whiten biasedvar eps 1e-08 tensor([-1.069045, -0.534522, 0.000000, 0.534522, 1.069045, 1.603567]) cand whiten unbiasvar eps 1e-08 tensor([-0.925820, -0.462910, 0.000000, 0.462910, 0.925820, 1.388730]) cand whiten biasedvar eps 0.0 tensor([-1.069045, -0.534522, 0.000000, 0.534522, 1.069045, 1.603567]) cand whiten unbiasvar eps 0.0 tensor([-0.925820, -0.462910, 0.000000, 0.462910, 0.925820, 1.388730]) entropy ref: 1.2569808959960938 masked_mean H: 1.2569807767868042 masked_sum H: 3.770942449569702 plain mean H: 1.2428988218307495 normalize: tensor([-0.888330, -0.565301, 0.080757, 1.372874]) (x-mean)/std_unb: tensor([-0.888330, -0.565301, 0.080757, 1.372874]) (x-mean)/std_bia: tensor([-1.025755, -0.652753, 0.093250, 1.585258]) (x-mean)/sqrt(var_b+eps): tensor([-1.025755, -0.652753, 0.093250, 1.585258]) dpo ls0 ref: 0.6782596707344055 cand: 0.6782596707344055 dpo ls.2 ref: 0.6842596530914307 cand: 0.6842596530914307 ipo ref: 22.089998245239258 cand (h-1/2b)^2: 22.089998245239258 bt ref: tensor([0.050000, 0.150000]) beta*(c-r): tensor([0.050000, 0.150000]) logsig(beta*(c-r)): tensor([-0.668460, -0.620957]) gae 1D: tensor([ 0.622257, 2.446058, -0.117673, -1.037425]) gae1d cand: tensor([ 0.622257, 2.446058, -0.117673, -1.037425]) lambda_returns 1D: tensor([ 0.893533, 1.173140, 0.385024, -0.619344]) lret cand adv+val: tensor([ 0.893533, 1.173140, 0.385024, -0.619344]) discounted 1D: tensor([2.750000, 3.500000, 3.000000]) cand: tensor([2.750000, 3.500000, 3.000000]) clip_fraction: 0.5 cand frac |ratio-1|>clip: 0.5 value_loss: 0.125 cand .5max: 0.125 cand max no .5: 0.25 huber: 0.125 smooth_l1: 0.125 smoothed_nll: 1.812943935394287 F.ce ls: 1.812943935394287 logprob_at_temp: tensor([-1.824993, -1.450230, -2.388001, -0.771483]) cand: tensor([-1.824993, -1.450230, -2.388001, -0.771483]) quantile_lower .25: 3.0 torch.quantile: 3.25 quantile lower interp: 3.0 median_reward even: 2.5 torch.median: 2.0 quantile.5: 2.5 [stdout] whiten shiftT: tensor([-1.069045, -0.534522, 0.000000, 0.534522, 1.069045, 1.603567]) whiten shiftF: tensor([1.930955, 2.465477, 3.000000, 3.534523, 4.069045, 4.603567]) masked_whiten shiftT: tensor([-1.069045, -0.534522, 0.000000, 0.000000, 0.000000, 1.603567]) mean 3.0 var_b 3.5 var_u 4.666666507720947 cand whiten biasedvar eps 1e-08 tensor([-1.069045, -0.534522, 0.000000, 0.534522, 1.069045, 1.603567]) cand whiten unbiasvar eps 1e-08 tensor([-0.925820, -0.462910, 0.000000, 0.462910, 0.925820, 1.388730]) cand whiten biasedvar eps 0.0 tensor([-1.069045, -0.534522, 0.000000, 0.534522, 1.069045, 1.603567]) cand whiten unbiasvar eps 0.0 tensor([-0.925820, -0.462910, 0.000000, 0.462910, 0.925820, 1.388730]) entropy ref: 1.2569808959960938 masked_mean H: 1.2569807767868042 masked_sum H: 3.770942449569702 plain mean H: 1.2428988218307495 normalize: tensor([-0.888330, -0.565301, 0.080757, 1.372874]) (x-mean)/std_unb: tensor([-0.888330, -0.565301, 0.080757, 1.372874]) (x-mean)/std_bia: tensor([-1.025755, -0.652753, 0.093250, 1.585258]) (x-mean)/sqrt(var_b+eps): tensor([-1.025755, -0.652753, 0.093250, 1.585258]) dpo ls0 ref: 0.6782596707344055 cand: 0.6782596707344055 dpo ls.2 ref: 0.6842596530914307 cand: 0.6842596530914307 ipo ref: 22.089998245239258 cand (h-1/2b)^2: 22.089998245239258 bt ref: tensor([0.050000, 0.150000]) beta*(c-r): tensor([0.050000, 0.150000]) logsig(beta*(c-r)): tensor([-0.668460, -0.620957]) gae 1D: tensor([ 0.622257, 2.446058, -0.117673, -1.037425]) gae1d cand: tensor([ 0.622257, 2.446058, -0.117673, -1.037425]) lambda_returns 1D: tensor([ 0.893533, 1.173140, 0.385024, -0.619344]) lret cand adv+val: tensor([ 0.893533, 1.173140, 0.385024, -0.619344]) discounted 1D: tensor([2.750000, 3.500000, 3.000000]) cand: tensor([2.750000, 3.500000, 3.000000]) clip_fraction: 0.5 cand frac |ratio-1|>clip: 0.5 value_loss: 0.125 cand .5max: 0.125 cand max no .5: 0.25 huber: 0.125 smooth_l1: 0.125 smoothed_nll: 1.812943935394287 F.ce ls: 1.812943935394287 logprob_at_temp: tensor([-1.824993, -1.450230, -2.388001, -0.771483]) cand: tensor([-1.824993, -1.450230, -2.388001, -0.771483]) quantile_lower .25: 3.0 torch.quantile: 3.25 quantile lower interp: 3.0 median_reward even: 2.5 torch.median: 2.0 quantile.5: 2.5
import torch, reforacle as R
import torch.nn.functional as F
torch.set_printoptions(precision=6,sci_mode=False)
G=lambda s:torch.Generator().manual_seed(s)
# grpo std/eps
rew=torch.tensor([1.,2,4, 0,10,-2.])
gt=R.grpo_advantages(rew,3,True); gf=R.grpo_advantages(rew,3,False)
print("grpo T",gt,"F",gf)
r2=rew.reshape(2,3); mu=r2.mean(1,keepdim=True)
for unb in [True,False]:
sd=r2.std(1,keepdim=True,unbiased=unb)
for eps in [1e-4,1e-6,1e-8,0.0]:
cand=((r2-mu)/(sd+eps)).flatten()
if torch.allclose(cand,gt,atol=1e-6): print(" MATCH unbiased",unb,"eps",eps)
# advantage_mean_std
adv=torch.tensor([1.,2,3,4,5]); m=torch.tensor([1.,1,0,1,0.])
ms=R.advantage_mean_std(adv,m); print("adv_mean_std",ms)
mv=(adv*m).sum()/m.sum()
for unb in [True,False]:
var=((adv-mv)**2*m).sum()/(m.sum()-(1 if unb else 0))
print(" cand std unb",unb, mv.item(), (var**0.5).item())
# top_p
probs=torch.tensor([[0.5,0.3,0.15,0.05],[0.4,0.4,0.1,0.1]])
print("top_p .8:",R.top_p_mask(probs,0.8))
print("top_p .7:",R.top_p_mask(probs,0.7))
print("top_p 1.0:",R.top_p_mask(probs,1.0))
# mode ties
print("mode [1,1,2,2]:",R.mode_label(torch.tensor([1,1,2,2])).item())
print("mode [3,3,1,1,2]:",R.mode_label(torch.tensor([3,3,1,1,2])).item())
# batched gae/discounted/lambda/cumulative
rew=torch.randn(3,4,generator=G(1)); val=torch.randn(3,4,generator=G(2))
print("gae (3,4) nv(4):",end=" ")
try: print(R.gae(rew,val,torch.randn(4,generator=G(3)),0.9,0.95).shape)
except Exception as e: print("ERR",e)
print("gae (3,4) nv(3):",end=" ")
try: print(R.gae(rew,val,torch.randn(3,generator=G(3)),0.9,0.95).shape)
except Exception as e: print("ERR",e)
dr=torch.tensor([[1.,2,3],[4,5,6.]])
print("discounted (2,3):",R.discounted_returns(dr,0.5)) # which axis
print("cumulative_max (2,3):",R.cumulative_max(torch.tensor([[1.,3,2],[5,1,9.]])))
# cross_entropy with ignore_index
lg=torch.randn(2,3,5,generator=G(4)); lb=torch.tensor([[0,1,-100],[2,-100,3]])
print("ce ignore:",R.cross_entropy(lg,lb,-100).item(), "F:",F.cross_entropy(lg.reshape(-1,5),lb.reshape(-1),ignore_index=-100).item())
# clipped_pg masked_mean vs sum
lp=torch.randn(2,3,generator=G(5)); olp=torch.randn(2,3,generator=G(6)); adv=torch.randn(2,3,generator=G(7)); m=torch.tensor([[1.,1,0],[1,0,1]])
r=R.clipped_pg_loss(lp,olp,adv,m,0.2,0.3)
ratio=torch.exp(lp-olp); un=ratio*adv; cl=torch.clamp(ratio,1-0.2,1+0.3)*adv; pg=-torch.min(un,cl)
print("clipped_pg ref",r.item(),"masked_mean",((pg*m).sum()/m.sum()).item())
Run third discovery
python3 d3.py 2>&1
grpo T tensor([-1.069044, -0.267261, 1.336305, -0.508000, 1.397001, -0.889001]) F tensor([-1.333333, -0.333333, 1.666667, -2.666667, 7.333333, -4.666667])
MATCH unbiased False eps 1e-06
MATCH unbiased False eps 1e-08
MATCH unbiased False eps 0.0
adv_mean_std tensor([2.333333, 1.247219])
cand std unb True 2.3333332538604736 1.5275253057479858
cand std unb False 2.3333332538604736 1.247219204902649
top_p .8: tensor([[0.625000, 0.375000, 0.000000, 0.000000],
[0.500000, 0.500000, 0.000000, 0.000000]])
top_p .7: tensor([[0.625000, 0.375000, 0.000000, 0.000000],
[0.500000, 0.500000, 0.000000, 0.000000]])
top_p 1.0: tensor([[0.500000, 0.300000, 0.150000, 0.050000],
[0.400000, 0.400000, 0.100000, 0.100000]])
mode [1,1,2,2]: 2
mode [3,3,1,1,2]: 3
gae (3,4) nv(4): torch.Size([3, 4])
gae (3,4) nv(3): ERR The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0
discounted (2,3): tensor([[3.000000, 4.500000, 6.000000],
[4.000000, 5.000000, 6.000000]])
cumulative_max (2,3): tensor([[1., 3., 2.],
[1., 3., 2.]])
ce ignore: 2.311739921569824 F: 2.311739921569824
clipped_pg ref 2.2903361320495605 masked_mean 2.2903361320495605
[stdout]
grpo T tensor([-1.069044, -0.267261, 1.336305, -0.508000, 1.397001, -0.889001]) F tensor([-1.333333, -0.333333, 1.666667, -2.666667, 7.333333, -4.666667])
MATCH unbiased False eps 1e-06
MATCH unbiased False eps 1e-08
MATCH unbiased False eps 0.0
adv_mean_std tensor([2.333333, 1.247219])
cand std unb True 2.3333332538604736 1.5275253057479858
cand std unb False 2.3333332538604736 1.247219204902649
top_p .8: tensor([[0.625000, 0.375000, 0.000000, 0.000000],
[0.500000, 0.500000, 0.000000, 0.000000]])
top_p .7: tensor([[0.625000, 0.375000, 0.000000, 0.000000],
[0.500000, 0.500000, 0.000000, 0.000000]])
top_p 1.0: tensor([[0.500000, 0.300000, 0.150000, 0.050000],
[0.400000, 0.400000, 0.100000, 0.100000]])
mode [1,1,2,2]: 2
mode [3,3,1,1,2]: 3
gae (3,4) nv(4): torch.Size([3, 4])
gae (3,4) nv(3): ERR The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0
discounted (2,3): tensor([[3.000000, 4.500000, 6.000000],
[4.000000, 5.000000, 6.000000]])
cumulative_max (2,3): tensor([[1., 3., 2.],
[1., 3., 2.]])
ce ignore: 2.311739921569824 F: 2.311739921569824
clipped_pg ref 2.2903361320495605 masked_mean 2.2903361320495605import torch, reforacle as R
import torch.nn.functional as F
torch.set_printoptions(precision=7,sci_mode=False)
G=lambda s:torch.Generator().manual_seed(s)
# grpo eps: tiny within-group std
rew=torch.tensor([1.0,1.0,1.0001, 5.,6.,7.])
gt=R.grpo_advantages(rew,3,True); print("grpo tiny",gt)
r2=rew.reshape(2,3);mu=r2.mean(1,keepdim=True);sd=r2.std(1,keepdim=True,unbiased=False)
for eps in [1e-4,1e-6,1e-8,1e-5]:
print(" eps",eps,((r2-mu)/(sd+eps)).flatten())
# cumulative_max dim
x=torch.tensor([[1.,3,2,5],[9,1,4,2.]])
print("cummax:",R.cumulative_max(x))
print("cummax dim0:",x.cummax(0).values)
print("cummax dim1:",x.cummax(1).values)
# mode: torch.mode
print("mode[3,3,1,1,2]:",torch.mode(torch.tensor([3,3,1,1,2])).values.item())
# ---------- COMPOSITES ----------
# dpo_sequence_loss
V=6
pc=torch.randn(2,4,V,generator=G(1)); pr=torch.randn(2,4,V,generator=G(2))
rc=torch.randn(2,4,V,generator=G(3)); rr=torch.randn(2,4,V,generator=G(4))
cl=torch.randint(0,V,(2,4),generator=G(5)); rl=torch.randint(0,V,(2,4),generator=G(6))
cm=(torch.rand(2,4,generator=G(7))>0.3).float(); rm=(torch.rand(2,4,generator=G(8))>0.3).float()
ref=R.dpo_sequence_loss(pc,pr,rc,rr,cl,rl,cm,rm,0.1,0.0)
def slp(lg,lb,m): return (F.log_softmax(lg,-1).gather(-1,lb.unsqueeze(-1)).squeeze(-1)*m).sum(-1)
pcl,prl,rcl,rrl=slp(pc,cl,cm),slp(pr,rl,rm),slp(rc,cl,cm),slp(rr,rl,rm)
h=0.1*((pcl-rcl)-(prl-rrl)); cand=(-F.logsigmoid(h)).mean()
print("dpo_seq ref",ref.item(),"cand",cand.item())
# rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high)
lg=torch.randn(6,4,V,generator=G(10)); olg=torch.randn(6,4,V,generator=G(11))
lb=torch.randint(0,V,(6,4),generator=G(12)); m=(torch.rand(6,4,generator=G(13))>0.3).float()
rew6=torch.randn(6,generator=G(14))
ref=R.rloo_objective(lg,olg,lb,m,rew6,3,0.2,0.3)
def tlp(l,lb): return F.log_softmax(l,-1).gather(-1,lb.unsqueeze(-1)).squeeze(-1)
logp=tlp(lg,lb); oldlp=tlp(olg,lb)
A=R.rloo_advantages(rew6,3) # per-sequence
ratio=torch.exp(logp-oldlp)
Ae=A.unsqueeze(-1)
pg=-torch.min(ratio*Ae, torch.clamp(ratio,1-0.2,1+0.3)*Ae)
print("rloo_obj ref",ref.item(),"cand tokmean",((pg*m).sum()/m.sum()).item())
# reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta)
rlg=torch.randn(6,4,V,generator=G(20))
adv=torch.randn(6,generator=G(21))
ref=R.reverse_kl_objective(lg,rlg,lb,m,adv,0.1)
logp=tlp(lg,lb); reflp=tlp(rlg,lb)
Ae=adv.unsqueeze(-1)
rkl=torch.exp(reflp-logp)-(reflp-logp)-1
pertok=-logp*Ae+0.1*rkl
print("rkl_obj ref",ref.item())
print(" cand pg+beta*rkl tokmean",((pertok*m).sum()/m.sum()).item())
print(" cand -A*logp only",(((-logp*Ae)*m).sum()/m.sum()).item())
print(" cand beta*rkl only",(((0.1*rkl)*m).sum()/m.sum()).item())
Probe composites and eps
python3 d4.py 2>&1
grpo tiny tensor([-0.6932461, -0.6932461, 1.3840162, -1.2247434, 0.0000000, 1.2247434])
eps 0.0001 tensor([-0.2268365, -0.2268365, 0.4528629, -1.2245948, 0.0000000, 1.2245948])
eps 1e-06 tensor([-0.6932461, -0.6932461, 1.3840162, -1.2247434, 0.0000000, 1.2247434])
eps 1e-08 tensor([-0.7077994, -0.7077994, 1.4130710, -1.2247448, 0.0000000, 1.2247448])
eps 1e-05 tensor([-0.5840701, -0.5840701, 1.1660544, -1.2247298, 0.0000000, 1.2247298])
cummax: tensor([[1., 3., 2., 5.],
[1., 3., 2., 5.]])
cummax dim0: tensor([[1., 3., 2., 5.],
[9., 3., 4., 5.]])
cummax dim1: tensor([[1., 3., 3., 5.],
[9., 9., 9., 9.]])
mode[3,3,1,1,2]: 1
dpo_seq ref 0.785122275352478 cand 0.785122275352478
rloo_obj ref 0.8786196112632751 cand tokmean 0.8786196112632751
rkl_obj ref -2.1043736934661865
cand pg+beta*rkl tokmean -2.1043736934661865
cand -A*logp only -2.2456839084625244
cand beta*rkl only 0.14130976796150208
[stdout]
grpo tiny tensor([-0.6932461, -0.6932461, 1.3840162, -1.2247434, 0.0000000, 1.2247434])
eps 0.0001 tensor([-0.2268365, -0.2268365, 0.4528629, -1.2245948, 0.0000000, 1.2245948])
eps 1e-06 tensor([-0.6932461, -0.6932461, 1.3840162, -1.2247434, 0.0000000, 1.2247434])
eps 1e-08 tensor([-0.7077994, -0.7077994, 1.4130710, -1.2247448, 0.0000000, 1.2247448])
eps 1e-05 tensor([-0.5840701, -0.5840701, 1.1660544, -1.2247298, 0.0000000, 1.2247298])
cummax: tensor([[1., 3., 2., 5.],
[1., 3., 2., 5.]])
cummax dim0: tensor([[1., 3., 2., 5.],
[9., 3., 4., 5.]])
cummax dim1: tensor([[1., 3., 3., 5.],
[9., 9., 9., 9.]])
mode[3,3,1,1,2]: 1
dpo_seq ref 0.785122275352478 cand 0.785122275352478
rloo_obj ref 0.8786196112632751 cand tokmean 0.8786196112632751
rkl_obj ref -2.1043736934661865
cand pg+beta*rkl tokmean -2.1043736934661865
cand -A*logp only -2.2456839084625244
cand beta*rkl only 0.14130976796150208import torch, reforacle as R
import torch.nn.functional as F
torch.set_printoptions(precision=7,sci_mode=False)
G=lambda s:torch.Generator().manual_seed(s)
def tlp(l,lb): return F.log_softmax(l,-1).gather(-1,lb.unsqueeze(-1)).squeeze(-1)
V=6
# cumulative_max 2D pattern
x=torch.tensor([[5.,1],[2,9],[3,3]])
print("cummax test:",R.cumulative_max(x))
print("1D:",R.cumulative_max(torch.tensor([2.,1,5,3,9,4])))
# grpo_objective
lg=torch.randn(6,4,V,generator=G(1)); olg=torch.randn(6,4,V,generator=G(2)); rlg=torch.randn(6,4,V,generator=G(3))
lb=torch.randint(0,V,(6,4),generator=G(4)); cm=(torch.rand(6,4,generator=G(5))>0.3).float()
rew=torch.randn(6,generator=G(6))
ref=R.grpo_objective(lg,olg,rlg,lb,cm,rew,3,0.05,0.2,0.3,True,"k3")
logp=tlp(lg,lb);oldlp=tlp(olg,lb);reflp=tlp(rlg,lb)
A=R.grpo_advantages(rew,3,True).unsqueeze(-1)
ratio=torch.exp(logp-oldlp)
pg=-torch.min(ratio*A,torch.clamp(ratio,1-0.2,1+0.3)*A)
kl=torch.exp(reflp-logp)-(reflp-logp)-1
pertok=pg+0.05*kl
print("grpo_obj ref",ref.item())
print(" cand tokmean(pg+beta*kl)",((pertok*cm).sum()/cm.sum()).item())
# alt: per-seq mean then mean
seq=(pertok*cm).sum(-1)/cm.sum(-1)
print(" cand seqmean",seq.mean().item())
# ppo_objective
T,B=5,3
rew=torch.randn(T,B,generator=G(10)); val=torch.randn(T,B,generator=G(11)); ov=torch.randn(T,B,generator=G(12))
logp=torch.randn(T,B,generator=G(13)); oldlp=torch.randn(T,B,generator=G(14)); nv=torch.randn(B,generator=G(15))
ref=R.ppo_objective(rew,val,ov,logp,oldlp,nv,0.99,0.95,0.2,0.3,0.2,0.5)
def gae(rew,val,nv,g,l):
T=rew.shape[0];adv=torch.zeros_like(rew);last=0.0
for t in reversed(range(T)):
nvv=nv if t==T-1 else val[t+1]
delta=rew[t]+g*nvv-val[t]; last=delta+g*l*last; adv[t]=last
return adv
adv=gae(rew,val,nv,0.99,0.95); ret=adv+val
ratio=torch.exp(logp-oldlp)
def clpg(adv):
pg=-torch.min(ratio*adv,torch.clamp(ratio,1-0.2,1+0.3)*adv); return pg.mean()
vpc=ov+torch.clamp(val-ov,-0.2,0.2); vl=0.5*torch.max((val-ret)**2,(vpc-ret)**2).mean()
print("ppo ref",ref.item())
print(" cand raw adv pg+0.5vf",(clpg(adv)+0.5*vl).item())
# whitened adv
w=(adv-adv.mean())*torch.rsqrt(adv.var(unbiased=False)+1e-8)
print(" cand whiten adv",(clpg(w)+0.5*vl).item())
wu=(adv-adv.mean())/(adv.std()+1e-8)
print(" cand normalize adv",(clpg(wu)+0.5*vl).item())
Probe grpo/ppo objectives
python3 d5.py 2>&1
cummax test: tensor([[5., 1.],
[5., 1.],
[5., 9.]])
1D: tensor([2., 2., 2., 5., 5., 9.])
grpo_obj ref 0.9544601440429688
cand tokmean(pg+beta*kl) 0.9544601440429688
cand seqmean 0.8327324986457825
ppo ref 2.5736236572265625
cand raw adv pg+0.5vf 3.679281711578369
cand whiten adv 2.5736236572265625
cand normalize adv 2.5292887687683105
[stdout]
cummax test: tensor([[5., 1.],
[5., 1.],
[5., 9.]])
1D: tensor([2., 2., 2., 5., 5., 9.])
grpo_obj ref 0.9544601440429688
cand tokmean(pg+beta*kl) 0.9544601440429688
cand seqmean 0.8327324986457825
ppo ref 2.5736236572265625
cand raw adv pg+0.5vf 3.679281711578369
cand whiten adv 2.5736236572265625
cand normalize adv 2.5292887687683105from __future__ import annotations
import math
import torch
import torch.nn.functional as F
# ----------------------------------------------------------------------------
# helpers
# ----------------------------------------------------------------------------
def _sum(x, axis):
return x.sum() if axis is None else x.sum(dim=axis)
def _token_logprobs(logits, labels):
lp = F.log_softmax(logits, dim=-1)
return lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
# ----------------------------------------------------------------------------
# masked primitives
# ----------------------------------------------------------------------------
def masked_mean(values, mask, axis=None):
s = _sum(values * mask, axis)
d = _sum(mask, axis)
return torch.where(d > 0, s / d, torch.zeros_like(s))
def masked_sum(values, mask, axis=None):
return _sum(values * mask, axis)
def logsumexp(x, axis):
return torch.logsumexp(x, dim=axis)
def log_softmax(x, axis):
return F.log_softmax(x, dim=axis)
def selective_logprobs(logits, labels, mask):
return (_token_logprobs(logits, labels) * mask).sum(dim=-1)
def token_logprobs(logits, labels):
return _token_logprobs(logits, labels)
def sequence_logprob(logits, labels, mask, length_normalize):
lp = _token_logprobs(logits, labels)
s = (lp * mask).sum(dim=-1)
if length_normalize:
return s / mask.sum(dim=-1)
return s
def entropy(logits, mask):
lp = F.log_softmax(logits, dim=-1)
p = lp.exp()
h = -(p * lp).sum(dim=-1)
return masked_mean(h, mask)
def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
logits = beta * ((pc - pr) - (rc - rr))
losses = -(1 - label_smoothing) * F.logsigmoid(logits) \
- label_smoothing * F.logsigmoid(-logits)
return losses.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):
g = rewards.reshape(-1, group_size)
mean = g.mean(dim=1, keepdim=True)
adv = g - mean
if scale_by_std:
std = g.std(dim=1, keepdim=True, unbiased=False)
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(next_value)
for t in reversed(range(T)):
nv = next_value if t == T - 1 else values[t + 1]
delta = rewards[t] + gamma * nv - 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 torch.exp(-d) + d - 1.0
raise ValueError(estimator)
def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
ratio = torch.exp(logp - old_logp)
unclipped = ratio * advantages
clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * advantages
loss = -torch.min(unclipped, clipped)
return masked_mean(loss, mask)
def value_loss(values, old_values, returns, clip):
if clip is None:
return 0.5 * ((values - returns) ** 2).mean()
vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)
loss = torch.max((values - returns) ** 2, (vpred_clipped - returns) ** 2)
return 0.5 * loss.mean()
def whiten(values, mask, shift_mean):
mean = masked_mean(values, mask)
var = masked_mean((values - mean) ** 2, mask)
whitened = (values - mean) * torch.rsqrt(var + 1e-8)
if not shift_mean:
whitened = whitened + mean
return whitened
def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
chosen_labels, rejected_labels, chosen_mask, rejected_mask,
beta, label_smoothing):
pc = (_token_logprobs(pc_logits, chosen_labels) * chosen_mask).sum(dim=-1)
pr = (_token_logprobs(pr_logits, rejected_labels) * rejected_mask).sum(dim=-1)
rc = (_token_logprobs(rc_logits, chosen_labels) * chosen_mask).sum(dim=-1)
rr = (_token_logprobs(rr_logits, rejected_labels) * rejected_mask).sum(dim=-1)
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)
old_logp = _token_logprobs(old_logits, labels)
ref_logp = _token_logprobs(ref_logits, labels)
adv = grpo_advantages(rewards, group_size, scale_by_std).unsqueeze(-1)
ratio = torch.exp(logp - old_logp)
unclipped = ratio * adv
clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv
pg = -torch.min(unclipped, clipped)
kl = kl_penalty(logp, ref_logp, kl_estimator)
per_token = pg + beta * kl
return masked_mean(per_token, completion_mask)
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
adv = whiten(adv, torch.ones_like(adv), True)
ratio = torch.exp(logp - old_logp)
unclipped = ratio * adv
clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv
pg = -torch.min(unclipped, clipped).mean()
vl = value_loss(values, old_values, returns, vf_clip)
return pg + vf_coef * vl
def rloo_advantages(rewards, group_size):
g = rewards.reshape(-1, group_size)
n = group_size
baseline = (g.sum(dim=1, keepdim=True) - g) / (n - 1)
adv = g - baseline
return adv.reshape(rewards.shape)
def reverse_kl(logp, ref_logp):
d = logp - ref_logp
return torch.exp(-d) + d - 1.0
def importance_ratio(logp, old_logp, clip):
ratio = torch.exp(logp - old_logp)
if clip is not None:
ratio = torch.clamp(ratio, 1.0 - clip, 1.0 + clip)
return ratio
def discounted_returns(rewards, gamma):
out = torch.zeros_like(rewards)
run = torch.zeros_like(rewards[0]) if rewards.dim() > 1 else 0.0
for t in reversed(range(rewards.shape[0])):
run = rewards[t] + gamma * run
out[t] = run
return out
def normalize(x, eps):
return (x - x.mean()) / (x.std() + eps)
def top_p_mask(probs, p):
sorted_probs, idx = torch.sort(probs, descending=True, dim=-1)
cum_before = torch.cumsum(sorted_probs, dim=-1) - sorted_probs
keep_sorted = cum_before < p
keep = torch.zeros_like(probs, dtype=torch.bool).scatter(-1, idx, keep_sorted)
out = probs * keep
return out / out.sum(dim=-1, keepdim=True)
def smoothed_nll(logits, labels, smoothing):
return F.cross_entropy(logits.reshape(-1, logits.shape[-1]),
labels.reshape(-1), label_smoothing=smoothing)
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)
old_logp = _token_logprobs(old_logits, labels)
adv = rloo_advantages(rewards, group_size).unsqueeze(-1)
ratio = torch.exp(logp - old_logp)
unclipped = ratio * adv
clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv
pg = -torch.min(unclipped, clipped)
return masked_mean(pg, mask)
def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
logp = _token_logprobs(logits, labels)
ref_logp = _token_logprobs(ref_logits, labels)
adv = advantages.unsqueeze(-1)
rkl = reverse_kl(logp, ref_logp)
per_token = -adv * logp + beta * rkl
return masked_mean(per_token, mask)
def cross_entropy(logits, labels, ignore_index):
return F.cross_entropy(logits.reshape(-1, logits.shape[-1]),
labels.reshape(-1), ignore_index=ignore_index)
def top_k_mask(logits, k):
_, idx = logits.topk(k, dim=-1)
return torch.zeros_like(logits, dtype=torch.bool).scatter(-1, idx, True)
def group_mean_baseline(rewards, group_size):
g = rewards.reshape(-1, group_size)
adv = g - g.mean(dim=1, keepdim=True)
return adv.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
return 0.5 * (torch.exp(d) + torch.exp(-d)) - 1.0
def huber_value_loss(values, returns, delta):
return F.smooth_l1_loss(values, returns, beta=delta)
def normalized_entropy(logits, mask):
return entropy(logits, mask) / math.log(logits.shape[-1])
def clip_fraction(logp, old_logp, clip):
ratio = torch.exp(logp - old_logp)
return ((ratio - 1.0).abs() > clip).float().mean()
def masked_whiten(values, mask, shift_mean):
mean = masked_mean(values, mask)
var = masked_mean((values - mean) ** 2, mask)
whitened = (values - mean) * torch.rsqrt(var + 1e-8)
whitened = whitened * mask
if not shift_mean:
whitened = whitened + mean
return whitened
def logprob_at_temperature(logits, labels, temperature):
return _token_logprobs(logits / temperature, labels)
def advantage_mean_std(advantages, mask):
mean = masked_mean(advantages, mask)
var = masked_mean((advantages - mean) ** 2, mask)
return mean, torch.sqrt(var)
def argmax_tokens(logits):
return logits.argmax(dim=-1)
def mode_label(labels):
return -torch.mode(-labels).values
def median_reward(rewards):
return torch.median(rewards)
def masked_max(values, mask):
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):
nz = mask != 0
any_nz = nz.any(dim=-1)
idx = nz.float().argmax(dim=-1)
return torch.where(any_nz, idx, torch.full_like(idx, -1))
def cumulative_max(x):
c = x.cummax(0).values
return torch.cat([c[:1], c[:-1]], dim=0)
def bucketize_reward(rewards, boundaries):
return torch.bucketize(rewards, boundaries)
import torch, importlib
import reforacle as R
import objectives as O
importlib.reload(O)
def close(a, b):
if isinstance(a, (tuple, list)):
return all(close(x, y) for x, y in zip(a, b))
a = torch.as_tensor(a); b = torch.as_tensor(b)
if a.shape != b.shape: return False
if a.dtype == torch.bool or not a.is_floating_point():
return torch.equal(a, b)
return torch.allclose(a.float(), b.float(), atol=1e-6, rtol=1e-5, equal_nan=True)
results = {}
def check(name, args, seed=0):
fr = getattr(R, name); fo = getattr(O, name)
try:
ref = fr(*args)
except Exception as e:
results.setdefault(name, []).append(("REF_ERR", str(e))); return
try:
out = fo(*args)
except Exception as e:
results.setdefault(name, []).append(("IMPL_ERR", str(e))); return
ok = close(ref, out)
results.setdefault(name, []).append(("OK" if ok else "MISMATCH", None))
if not ok:
print(f"MISMATCH {name}:")
print(" ref ", torch.as_tensor(ref).flatten()[:8].tolist() if not isinstance(ref,tuple) else [torch.as_tensor(r).tolist() for r in ref])
print(" out ", torch.as_tensor(out).flatten()[:8].tolist() if not isinstance(out,tuple) else [torch.as_tensor(r).tolist() for r in out])
def G(s): return torch.Generator().manual_seed(s)
for s in range(4):
rn = lambda *sh: torch.randn(*sh, generator=G(s*100+hash(sh)%50))
g = G(s)
def R_(*sh): return torch.randn(*sh, generator=g)
def U(*sh): return torch.rand(*sh, generator=g)
V = 7
v = R_(4, 5); m = (U(4, 5) > 0.4).float()
check("masked_mean", (v, m))
check("masked_mean", (v, m, 1))
check("masked_sum", (v, m))
check("masked_sum", (v, m, 0))
check("masked_max", (v, m))
check("logsumexp", (R_(3, 4), 1))
check("log_softmax", (R_(3, 4), 0))
check("whiten", (R_(10), (U(10) > .3).float(), True))
check("whiten", (R_(10), (U(10) > .3).float(), False))
check("masked_whiten", (v, m, True))
check("masked_whiten", (v, m, False))
logits = R_(2, 6, V); labels = torch.randint(0, V, (2, 6), generator=g); sm = (U(2, 6) > .3).float()
check("token_logprobs", (logits, labels))
check("selective_logprobs", (logits, labels, sm))
check("sequence_logprob", (logits, labels, sm, False))
check("sequence_logprob", (logits, labels, sm, True))
check("entropy", (logits, sm))
check("normalized_entropy", (logits, sm))
check("cross_entropy", (logits, labels, -100))
lb2 = labels.clone(); lb2[0, 0] = -100
check("cross_entropy", (logits, lb2, -100))
check("smoothed_nll", (logits, labels, 0.1))
check("logprob_at_temperature", (logits, labels, 0.7))
check("argmax_tokens", (logits,))
pc, pr, rc, rr = R_(5), R_(5), R_(5), R_(5)
check("dpo_loss", (pc, pr, rc, rr, 0.1, 0.0))
check("dpo_loss", (pc, pr, rc, rr, 0.2, 0.15))
check("ipo_loss", (pc, pr, rc, rr, 0.1))
check("bradley_terry_logit", (R_(5), R_(5), 0.3))
rew = R_(9)
check("grpo_advantages", (rew, 3, True))
check("grpo_advantages", (rew, 3, False))
check("rloo_advantages", (rew, 3))
check("group_mean_baseline", (rew, 3))
check("advantage_mean_std", (R_(8), (U(8) > .3).float()))
rw = R_(5, 3); vl = R_(5, 3); nv = R_(3)
check("gae", (rw, vl, nv, 0.99, 0.95))
check("lambda_returns", (rw, vl, nv, 0.99, 0.95))
check("discounted_returns", (R_(6), 0.97))
check("discounted_returns", (rw, 0.9))
lp, rlp = R_(3, 4), R_(3, 4)
for est in ("k1", "k2", "k3"):
check("kl_penalty", (lp, rlp, est))
check("reverse_kl", (lp, rlp))
check("symmetric_kl", (lp, rlp))
olp = R_(3, 4); adv = R_(3, 4); mm = (U(3, 4) > .3).float()
check("clipped_pg_loss", (lp, olp, adv, mm, 0.2, 0.3))
check("importance_ratio", (lp, olp, None))
check("importance_ratio", (lp, olp, 0.2))
check("clip_fraction", (lp, olp, 0.2))
val, ov, ret = R_(3, 4), R_(3, 4), R_(3, 4)
check("value_loss", (val, ov, ret, 0.2))
check("huber_value_loss", (val, ret, 1.0))
check("normalize", (R_(6), 1e-8))
check("top_p_mask", (torch.softmax(R_(3, 6), -1), 0.8))
check("top_k_mask", (R_(3, 6), 3))
check("mode_label", (torch.randint(0, 4, (12,), generator=g),))
check("median_reward", (R_(7),))
check("quantile_lower", (R_(11), 0.3))
check("pad_mask_from_lengths", (torch.randint(0, 5, (4,), generator=g), 5))
check("first_nonzero_index", ((U(4, 5) > .6).float(),))
check("cumulative_max", (R_(8),))
check("cumulative_max", (R_(4, 3),))
check("bucketize_reward", (R_(6), torch.tensor([-0.5, 0.0, 0.5, 1.0])))
# composites
pcl = R_(2, 4, V); prl = R_(2, 4, V); rcl = R_(2, 4, V); rrl = R_(2, 4, V)
clab = torch.randint(0, V, (2, 4), generator=g); rlab = torch.randint(0, V, (2, 4), generator=g)
cmk = (U(2, 4) > .3).float(); rmk = (U(2, 4) > .3).float()
check("dpo_sequence_loss", (pcl, prl, rcl, rrl, clab, rlab, cmk, rmk, 0.1, 0.1))
lg = R_(6, 4, V); olg = R_(6, 4, V); rlg = R_(6, 4, V)
lab = torch.randint(0, V, (6, 4), generator=g); cm = (U(6, 4) > .3).float()
rew6 = R_(6)
check("grpo_objective", (lg, olg, rlg, lab, cm, rew6, 3, 0.05, 0.2, 0.3, True, "k3"))
check("grpo_objective", (lg, olg, rlg, lab, cm, rew6, 2, 0.04, 0.2, 0.2, False, "k1"))
check("rloo_objective", (lg, olg, lab, cm, rew6, 3, 0.2, 0.3))
check("reverse_kl_objective", (lg, rlg, lab, cm, R_(6), 0.1))
T, B = 5, 3
check("ppo_objective", (R_(T, B), R_(T, B), R_(T, B), R_(T, B), R_(T, B), R_(B), 0.99, 0.95, 0.2, 0.3, 0.2, 0.5))
fails = {n: v for n, v in results.items() if any(s[0] != "OK" for s in v)}
print("\n==== SUMMARY ====")
print(f"functions tested: {len(results)}")
if fails:
for n, v in fails.items():
st = set(s[0] for s in v)
print(f"FAIL {n}: {st} {[s[1] for s in v if s[1]][:1]}")
else:
print("ALL PASS")
Run comparison harness
python3 harness.py 2>&1 | tail -40
==== SUMMARY ==== functions tested: 49 ALL PASS [stdout] ==== SUMMARY ==== functions tested: 49 ALL PASS
import torch
import reforacle as R
import objectives as O
def cmp(name, args, tol=1e-5):
fr, fo = getattr(R, name), getattr(O, name)
try: ref = fr(*args)
except Exception as e: ref = ("ERR", type(e).__name__)
try: out = fo(*args)
except Exception as e: out = ("ERR", type(e).__name__)
def norm(z):
if isinstance(z, tuple) and z and z[0]=="ERR": return z
if isinstance(z,(tuple,list)): return tuple(torch.as_tensor(x) for x in z)
return torch.as_tensor(z)
r,o=norm(ref),norm(out)
ok=False
try:
if isinstance(r,tuple) and r and r[0]=="ERR":
ok = (o[0]=="ERR")
elif isinstance(r,tuple):
ok=all(torch.allclose(a.float(),b.float(),atol=tol,rtol=tol,equal_nan=True) for a,b in zip(r,o))
else:
ok = r.shape==o.shape and (torch.equal(r,o) if (r.dtype==torch.bool or not r.is_floating_point()) else torch.allclose(r.float(),o.float(),atol=tol,rtol=tol,equal_nan=True))
except Exception as e:
ok=False
print(("ok " if ok else "FAIL")+f" {name}", "" if ok else f"\n ref={ref}\n out={out}")
g=torch.Generator().manual_seed(42)
def rn(*s): return torch.randn(*s,generator=g)
def ru(*s): return torch.rand(*s,generator=g)
# empty / all-zero masks
z=torch.zeros(5)
cmp("masked_mean",(rn(5),z))
cmp("masked_sum",(rn(5),z))
cmp("entropy",(rn(3,4),torch.zeros(3)))
cmp("advantage_mean_std",(rn(5),z))
cmp("whiten",(rn(6),torch.zeros(6),True))
cmp("masked_whiten",(rn(6),torch.zeros(6),True))
cmp("masked_max",(rn(5),torch.tensor([0.,0,1,0,0])))
# 1D / 2D token functions
cmp("token_logprobs",(rn(5,8),torch.randint(0,8,(5,),generator=g)))
cmp("selective_logprobs",(rn(5,8),torch.randint(0,8,(5,),generator=g),(ru(5,8)>.3).float()))
cmp("logsumexp",(rn(3,4,5),2))
cmp("log_softmax",(rn(3,4,5),-1))
cmp("cross_entropy",(rn(4,7),torch.randint(0,7,(4,),generator=g),-100))
cmp("smoothed_nll",(rn(4,7),torch.randint(0,7,(4,),generator=g),0.2))
cmp("entropy",(rn(3,4,9),(ru(3,4)>.3).float()))
cmp("normalized_entropy",(rn(3,4,9),(ru(3,4)>.3).float()))
# gae 1D & 3D-ish
cmp("gae",(rn(6),rn(6),torch.tensor(0.3),0.99,0.95))
cmp("lambda_returns",(rn(6),rn(6),torch.tensor(0.3),0.99,0.95))
cmp("discounted_returns",(rn(5),1.0))
# grpo group sizes / negative
cmp("grpo_advantages",(rn(12),4,True))
cmp("grpo_advantages",(rn(12),6,True))
cmp("rloo_advantages",(rn(12),4))
# quantile / median even & odd, ties
cmp("median_reward",(torch.tensor([3.,1,2,4]),))
cmp("quantile_lower",(torch.tensor([1.,2,3,4,5,6]),0.5))
cmp("quantile_lower",(rn(20),0.9))
cmp("mode_label",(torch.tensor([5,5,2,2,7,7,1]),))
cmp("mode_label",(torch.tensor([0,0,0,3,3]),))
# top_p edges
p=torch.softmax(rn(4,10),-1)
cmp("top_p_mask",(p,0.0))
cmp("top_p_mask",(p,1.0))
cmp("top_p_mask",(p,0.5))
cmp("top_k_mask",(rn(4,10),1))
cmp("top_k_mask",(rn(4,10),10))
# value_loss no clip
cmp("value_loss",(rn(3,4),rn(3,4),rn(3,4),None))
cmp("huber_value_loss",(rn(5),rn(5),0.5))
cmp("importance_ratio",(rn(4),rn(4),0.1))
cmp("clip_fraction",(rn(4),rn(4),0.0))
cmp("first_nonzero_index",(torch.zeros(3,4),))
cmp("cumulative_max",(rn(5,2),))
cmp("bucketize_reward",(torch.tensor([-1.,0.0,0.5,2.0]),torch.tensor([0.,0.5,1.0])))
cmp("pad_mask_from_lengths",(torch.tensor([0,5,3]),5))
Run edge case tests
python3 edge.py 2>&1
ok masked_mean
ok masked_sum
ok entropy
FAIL advantage_mean_std
ref=tensor([0., 0.])
out=(tensor(0.), tensor(0.))
ok whiten
ok masked_whiten
ok masked_max
ok token_logprobs
ok selective_logprobs
ok logsumexp
ok log_softmax
ok cross_entropy
ok smoothed_nll
ok entropy
ok normalized_entropy
ok gae
ok lambda_returns
ok discounted_returns
ok grpo_advantages
ok grpo_advantages
ok rloo_advantages
FAIL median_reward
ref=2.5
out=2.0
ok quantile_lower
ok quantile_lower
ok mode_label
ok mode_label
FAIL top_p_mask
ref=tensor([[0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],
[0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
[0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
[0., 0., 1., 0., 0., 0., 0., 0., 0., 0.]])
out=tensor([[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]])
ok top_p_mask
ok top_p_mask
ok top_k_mask
ok top_k_mask
FAIL value_loss
ref=('ERR', 'TypeError')
out=0.6920984387397766
FAIL huber_value_loss
ref=0.4619688391685486
out=0.9239376783370972
ok importance_ratio
ok clip_fraction
ok first_nonzero_index
ok cumulative_max
FAIL bucketize_reward
ref=tensor([0, 1, 2, 3])
out=tensor([0, 0, 1, 3])
ok pad_mask_from_lengths
[stdout]
ok masked_mean
ok masked_sum
ok entropy
FAIL advantage_mean_std
ref=tensor([0., 0.])
out=(tensor(0.), tensor(0.))
ok whiten
ok masked_whiten
ok masked_max
ok token_logprobs
ok selective_logprobs
ok logsumexp
ok log_softmax
ok cross_entropy
ok smoothed_nll
ok entropy
ok normalized_entropy
ok gae
ok lambda_returns
ok discounted_returns
ok grpo_advantages
ok grpo_advantages
ok rloo_advantages
FAIL median_reward
ref=2.5
out=2.0
ok quantile_lower
ok quantile_lower
ok mode_label
ok mode_label
FAIL top_p_mask
ref=tensor([[0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],
[0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
[0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
[0., 0., 1., 0., 0., 0., 0., 0., 0., 0.]])
out=tensor([[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]])
ok top_p_mask
ok top_p_mask
ok top_k_mask
ok top_k_mask
FAIL value_loss
ref=('ERR', 'TypeError')
out=0.6920984387397766
FAIL huber_value_loss
ref=0.4619688391685486
out=0.9239376783370972
ok importance_ratio
ok clip_fraction
ok first_nonzero_index
ok cumulative_max
FAIL bucketize_reward
ref=tensor([0, 1, 2, 3])
out=tensor([0, 0, 1, 3])
ok pad_mask_from_lengths/app/objectives.py
def value_loss(values, old_values, returns, clip):
if clip is None:
return 0.5 * ((values - returns) ** 2).mean()
vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)def value_loss(values, old_values, returns, clip):
vpred_clipped = old_values + torch.clamp(values - old_values, -clip, clip)/app/objectives.py
keep_sorted = cum_before < p
keep = torch.zeros_like(probs, dtype=torch.bool).scatter(-1, idx, keep_sorted) keep_sorted = cum_before < p
keep_sorted[..., 0] = True
keep = torch.zeros_like(probs, dtype=torch.bool).scatter(-1, idx, keep_sorted)/app/objectives.py
def huber_value_loss(values, returns, delta):
return F.smooth_l1_loss(values, returns, beta=delta)def huber_value_loss(values, returns, delta):
return F.huber_loss(values, returns, delta=delta)/app/objectives.py
var = masked_mean((advantages - mean) ** 2, mask)
return mean, torch.sqrt(var) var = masked_mean((advantages - mean) ** 2, mask)
return torch.stack([mean, torch.sqrt(var)])/app/objectives.py
def median_reward(rewards):
return torch.median(rewards)def median_reward(rewards):
return torch.quantile(rewards, 0.5)/app/objectives.py
return torch.bucketize(rewards, boundaries)
return torch.bucketize(rewards, boundaries, right=True)
Re-run edge and main harness
python3 edge.py 2>&1 | grep -v '^ok' ; echo "=== MAIN ===" && python3 harness.py 2>&1 | tail -5
=== MAIN === ==== SUMMARY ==== functions tested: 49 ALL PASS [stdout] === MAIN === ==== SUMMARY ==== functions tested: 49 ALL PASS
Broad stress test extra shapes
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(7)
def rn(*s): return torch.randn(*s,generator=g)
def ru(*s): return torch.rand(*s,generator=g)
def ck(n,a,tol=1e-5):
r=getattr(R,n)(*a); o=getattr(O,n)(*a)
r=torch.as_tensor(r); o=torch.as_tensor(o)
ok = r.shape==o.shape and (torch.equal(r,o) if (r.dtype==torch.bool or not r.is_floating_point()) else torch.allclose(r.float(),o.float(),atol=tol,rtol=tol,equal_nan=True))
print(('ok ' if ok else 'FAIL'),n, '' if ok else (r.flatten()[:5].tolist(),o.flatten()[:5].tolist()))
# 3D masked ops, neg axis
v=rn(2,3,4); m=(ru(2,3,4)>.4).float()
ck('masked_mean',(v,m,-1)); ck('masked_mean',(v,m,2)); ck('masked_sum',(v,m,-1))
ck('advantage_mean_std',(rn(9),(ru(9)>.3).float()))
ck('huber_value_loss',(rn(6),rn(6),0.3)); ck('huber_value_loss',(rn(6),rn(6),2.0))
ck('value_loss',(rn(4),rn(4),rn(4),0.5))
ck('median_reward',(rn(8),)); ck('median_reward',(rn(9),))
ck('bucketize_reward',(rn(10),torch.tensor([-1.,-.3,.2,.8])))
ck('top_p_mask',(torch.softmax(rn(5,12),-1),0.3))
ck('grpo_advantages',(rn(20),5,True)); ck('rloo_advantages',(rn(20),5))
ck('kl_penalty',(rn(4,5),rn(4,5),'k3')); ck('symmetric_kl',(rn(4,5),rn(4,5)))
ck('cumulative_max',(rn(7),)); ck('cumulative_max',(rn(3,4),))
ck('gae',(rn(8,2),rn(8,2),rn(2),0.95,0.9))
ck('ppo_objective',(rn(6,4),rn(6,4),rn(6,4),rn(6,4),rn(6,4),rn(4),0.99,0.9,0.15,0.25,0.3,0.7))
ck('grpo_objective',(rn(8,5,6),rn(8,5,6),rn(8,5,6),torch.randint(0,6,(8,5),generator=g),(ru(8,5)>.3).float(),rn(8),4,0.1,0.15,0.25,True,'k2'))
" 2>&1ok masked_mean ok masked_mean ok masked_sum ok advantage_mean_std ok huber_value_loss ok huber_value_loss ok value_loss ok median_reward ok median_reward ok bucketize_reward ok top_p_mask ok grpo_advantages ok rloo_advantages ok kl_penalty ok symmetric_kl ok cumulative_max ok cumulative_max ok gae ok ppo_objective ok grpo_objective [stdout] ok masked_mean ok masked_mean ok masked_sum ok advantage_mean_std ok huber_value_loss ok huber_value_loss ok value_loss ok median_reward ok median_reward ok bucketize_reward ok top_p_mask ok grpo_advantages ok rloo_advantages ok kl_penalty ok symmetric_kl ok cumulative_max ok cumulative_max ok gae ok ppo_objective ok grpo_objective
Verify standalone, no reforacle dependency
cd /tmp && python3 -c "
import ast
src=open('/app/objectives.py').read()
tree=ast.parse(src)
imps=[n.names[0].name for n in ast.walk(tree) if isinstance(n,(ast.Import,))]
imps+= [n.module for n in ast.walk(tree) if isinstance(n,ast.ImportFrom)]
print('imports:',imps)
assert not any('reforacle' in (i or '') for i in imps), 'reforacle referenced!'
assert 'reforacle' not in src
import sys; sys.path.insert(0,'/app')
import objectives # imports fine with only torch
funcs=[k for k in dir(objectives) if not k.startswith('_') and callable(getattr(objectives,k))]
print('callable funcs:', len([f for f in funcs if f not in ('torch','math','F')]))
print('OK standalone')
" && bash /app/timer.shimports: ['math', 'torch', 'torch.nn.functional', '__future__'] callable funcs: 49 OK standalone Remaining time (hours:minutes): 1:46 Shell cwd was reset to /app [stdout] imports: ['math', 'torch', 'torch.nn.functional', '__future__'] callable funcs: 49 OK standalone Remaining time (hours:minutes): 1:46 [stderr] Shell cwd was reset to /app
No file changes, the agent left the workspace untouched.
=== SyncVals verifier (post-training objectives, sealed oracle) === FUNCTIONS PASSED: 42/49 (partial_score=0.8571) 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_02a684ccebd449c2. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.
Trial trial_02a684ccebd449c2 · verifier authoritative; classifier explanatory.