SyncValsverifier → artifact → classifier → verdict
SyncVals · Trajectory

reverse-engineer-objectives

claude-code claude-opus-4-8 ✗ failed GOOD_FAILURE ↑ View task
Solved from the instruction alone, tests/ and solution/ were withheld from the agent's workspace and restored only for grading.
Reward = tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it.
Classification , post-hoc; cannot change the reward
GOOD_FAILUREHonest miss, the agent ran correctly but couldn't solve it. Expected for a hard task; the task is sound.
SubtypeIncomplete Exploration
EvidenceAgent achieved 49/49 local function tests (step 40: '49/49 functions pass') against the oracle, but verifier shows only 44/49 functions passing on held-out inputs (reward=0.0, requires all 49 to pass). Test output: 'FUNCTIONS PASSED: 44/49 (partial_score=0.898) REWARD = 0 (1 iff ALL 49 functions match the sealed reference exactly)'. Agent's probing methodology was sound, systematically discovering conventions for entropy, KL estimators, normalization, etc., but 5 functions had undiscovered edge-case behaviors that only manifested on hidden test inputs.
Root causeThe agent successfully reverse-engineered and implemented 44 of 49 post-training objectives by probing the oracle, but the 5 failing functions expose edge cases or numerical conventions that weren't exercised during the agent's probing session. The task explicitly states 'Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on,' making this an expected failure mode for reverse-engineering tasks.
RecommendationN/A - task is fine
Trajectory
Tool-by-tool agent trajectory
42 tool calls · 4 tool types · 44 steps
Reimplement the post-training objectives library in `/app/objectives.py` so that every function reproduces the sealed reference module `reforacle` exactly. `reforacle` is importable and callable (`import reforacle`) but its source cannot be read. It is the only ground truth; there is no written specification of what any function computes or of the conventions it follows. Recover its behavior by calling it. `/app/objectives.py` contains forty-nine function stubs. Keep every signature exactly as given and implement every function. The names indicate the standard post-training quantity each is meant to produce, but the reference's exact definition is whatever `reforacle` returns , not what you assume. Determine each definition by probing. Grading compares your implementation to `reforacle` on held-out inputs you do not see, function by function, requiring exact numerical agreement (primitives within `atol=1e-6, rtol=1e-5`; composite objectives within `atol=1e-5, rtol=1e-4`). Each function passes only if it agrees on every graded input, and the task passes only if all forty-nine functions pass. Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading runs your `/app/objectives.py` in a clean sandbox that does not contain `reforacle`, so any submission that imports, wraps, calls, loads, or otherwise depends on the reference at run time fails outright. Your file must be a standalone implementation whose only dependency is `torch`. Do not read, decompile, disassemble, or otherwise extract the compiled reference; recover its behavior only by calling it and comparing. `python3 /app/probe.py` shows the calling pattern; `bash /app/timer.sh` reports remaining budget.

/app/probe.py

contents
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

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

Remaining time (hours:minutes):
1:59

[stdout]
Remaining time (hours:minutes):
1:59

/app/objectives.py

contents
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

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

Remaining time (hours:minutes):
1:59

[stdout]
Remaining time (hours:minutes):
1:59

Check remaining budget

shell
bash /app/timer.sh
output
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

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

Remaining time (hours:minutes):
1:59

[stdout]
Remaining time (hours:minutes):
1:59

List reforacle functions and signatures

shell
python3 -c "
import reforacle, inspect
fns = [x for x in dir(reforacle) if not x.startswith('_')]
print(len(fns))
for f in fns:
    obj = getattr(reforacle, f)
    try:
        print(f, inspect.signature(obj))
    except Exception as e:
        print(f, 'NOSIG', type(obj))
"
output
51
F NOSIG <class 'module'>
advantage_mean_std (advantages, mask)
argmax_tokens (logits)
bradley_terry_logit (chosen_reward, rejected_reward, beta)
bucketize_reward (rewards, boundaries)
clip_fraction (logp, old_logp, clip)
clipped_pg_loss (logp, old_logp, advantages, mask, clip_low, clip_high)
cross_entropy (logits, labels, ignore_index)
cumulative_max (x)
discounted_returns (rewards, gamma)
dpo_loss (pc, pr, rc, rr, beta, label_smoothing)
dpo_sequence_loss (pc_logits, pr_logits, rc_logits, rr_logits, chosen_labels, rejected_labels, chosen_mask, rejected_mask, beta, label_smoothing)
entropy (logits, mask)
first_nonzero_index (mask)
gae (rewards, values, next_value, gamma, lam)
group_mean_baseline (rewards, group_size)
grpo_advantages (rewards, group_size, scale_by_std)
grpo_objective (logits, old_logits, ref_logits, labels, completion_mask, rewards, group_size, beta, clip_low, clip_high, scale_by_std, kl_estimator)
huber_value_loss (values, returns, delta)
importance_ratio (logp, old_logp, clip)
ipo_loss (pc, pr, rc, rr, beta)
kl_penalty (logp, ref_logp, estimator)
lambda_returns (rewards, values, next_value, gamma, lam)
log_softmax (x, axis)
logprob_at_temperature (logits, labels, temperature)
logsumexp (x, axis)
masked_max (values, mask)
masked_mean (values, mask, axis=None)
masked_sum (values, mask, axis=None)
masked_whiten (values, mask, shift_mean)
median_reward (rewards)
mode_label (labels)
normalize (x, eps)
normalized_entropy (logits, mask)
pad_mask_from_lengths (lengths, max_len)
ppo_objective (rewards, values, old_values, logp, old_logp, next_value, gamma, lam, clip_low, clip_high, vf_clip, vf_coef)
quantile_lower (x, q)
reverse_kl (logp, ref_logp)
reverse_kl_objective (logits, ref_logits, labels, mask, advantages, beta)
rloo_advantages (rewards, group_size)
rloo_objective (logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high)
selective_logprobs (logits, labels, mask)
sequence_logprob (logits, labels, mask, length_normalize)
smoothed_nll (logits, labels, smoothing)
symmetric_kl (logp, ref_logp)
token_logprobs (logits, labels)
top_k_mask (logits, k)
top_p_mask (probs, p)
torch NOSIG <class 'module'>
value_loss (values, old_values, returns, clip)
whiten (values, mask, shift_mean)

[stdout]
51
F NOSIG <class 'module'>
advantage_mean_std (advantages, mask)
argmax_tokens (logits)
bradley_terry_logit (chosen_reward, rejected_reward, beta)
bucketize_reward (rewards, boundaries)
clip_fraction (logp, old_logp, clip)
clipped_pg_loss (logp, old_logp, advantages, mask, clip_low, clip_high)
cross_entropy (logits, labels, ignore_index)
cumulative_max (x)
discounted_returns (rewards, gamma)
dpo_loss (pc, pr, rc, rr, beta, label_smoothing)
dpo_sequence_loss (pc_logits, pr_logits, rc_logits, rr_logits, chosen_labels, rejected_labels, chosen_mask, rejected_mask, beta, label_smoothing)
entropy (logits, mask)
first_nonzero_index (mask)
gae (rewards, values, next_value, gamma, lam)
group_mean_baseline (rewards, group_size)
grpo_advantages (rewards, group_size, scale_by_std)
grpo_objective (logits, old_logits, ref_logits, labels, completion_mask, rewards, group_size, beta, clip_low, clip_high, scale_by_std, kl_estimator)
huber_value_loss (values, returns, delta)
importance_ratio (logp, old_logp, clip)
ipo_loss (pc, pr, rc, rr, beta)
kl_penalty (logp, ref_logp, estimator)
lambda_returns (rewards, values, next_value, gamma, lam)
log_softmax (x, axis)
logprob_at_temperature (logits, labels, temperature)
logsumexp (x, axis)
masked_max (values, mask)
masked_mean (values, mask, axis=None)
masked_sum (values, mask, axis=None)
masked_whiten (values, mask, shift_mean)
median_reward (rewards)
mode_label (labels)
normalize (x, eps)
normalized_entropy (logits, mask)
pad_mask_from_lengths (lengths, max_len)
ppo_objective (rewards, values, old_values, logp, old_logp, next_value, gamma, lam, clip_low, clip_high, vf_clip, vf_coef)
quantile_lower (x, q)
reverse_kl (logp, ref_logp)
reverse_kl_objective (logits, ref_logits, labels, mask, advantages, beta)
rloo_advantages (rewards, group_size)
rloo_objective (logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high)
selective_logprobs (logits, labels, mask)
sequence_logprob (logits, labels, mask, length_normalize)
smoothed_nll (logits, labels, smoothing)
symmetric_kl (logp, ref_logp)
token_logprobs (logits, labels)
top_k_mask (logits, k)
top_p_mask (probs, p)
torch NOSIG <class 'module'>
value_loss (values, old_values, returns, clip)
whiten (values, mask, shift_mean)

Run the probe example

shell
python3 probe.py
output
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.

Probe entropy and token_logprobs conventions

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(1)
def rn(*s): return torch.randn(*s, generator=g)

# entropy
logits = rn(3,5)
mask = (torch.rand(3,generator=g)>0.3).float()
print('entropy', R.entropy(logits, mask))
p = torch.softmax(logits,-1)
ent = -(p*torch.log_softmax(logits,-1)).sum(-1)
print('  per-token ent', ent)
print('  masked_mean', (ent*mask).sum()/mask.sum())
print('  plain mean', ent.mean())
print('normalized_entropy', R.normalized_entropy(logits, mask))
print('  /logV', ((ent*mask).sum()/mask.sum())/torch.log(torch.tensor(5.0)))

# token_logprobs
logits2 = rn(2,4,6); labels=torch.randint(0,6,(2,4),generator=g)
tl = R.token_logprobs(logits2, labels)
print('token_logprobs shape', tl.shape)
print(tl[0,0], torch.log_softmax(logits2,-1)[0,0,labels[0,0]])
"
output
entropy tensor(1.5408)
  per-token ent tensor([1.5408, 1.4146, 1.4812])
  masked_mean tensor(1.5408)
  plain mean tensor(1.4789)
normalized_entropy tensor(0.9573)
  /logV tensor(0.9573)
token_logprobs shape torch.Size([2, 4])
tensor(-1.3877) tensor(-1.3877)

[stdout]
entropy tensor(1.5408)
  per-token ent tensor([1.5408, 1.4146, 1.4812])
  masked_mean tensor(1.5408)
  plain mean tensor(1.4789)
normalized_entropy tensor(0.9573)
  /logV tensor(0.9573)
token_logprobs shape torch.Size([2, 4])
tensor(-1.3877) tensor(-1.3877)

Probe selective_logprobs and sequence_logprob

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(2)
def rn(*s): return torch.randn(*s, generator=g)
logits = rn(2,4,6); labels=torch.randint(0,6,(2,4),generator=g)
mask=(torch.rand(2,4,generator=g)>0.3).float()
tl = torch.log_softmax(logits,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
sl=R.selective_logprobs(logits,labels,mask)
print('selective shape', sl.shape)
print('sel', sl)
print('tl*mask', tl*mask)
for ln in [False,True]:
    sq=R.sequence_logprob(logits,labels,mask,ln)
    print('seqlogprob ln=',ln, sq)
print('  sum tl*mask', (tl*mask).sum(-1))
print('  /masksum', (tl*mask).sum(-1)/mask.sum(-1))
"
output
selective shape torch.Size([2])
sel tensor([-9.9312, -7.4454])
tl*mask tensor([[-3.2451, -2.5404, -1.9318, -2.2140],
        [-3.0121, -0.0000, -3.2752, -1.1581]])
seqlogprob ln= False tensor([-9.9312, -7.4454])
seqlogprob ln= True tensor([-2.4828, -2.4818])
  sum tl*mask tensor([-9.9312, -7.4454])
  /masksum tensor([-2.4828, -2.4818])

[stdout]
selective shape torch.Size([2])
sel tensor([-9.9312, -7.4454])
tl*mask tensor([[-3.2451, -2.5404, -1.9318, -2.2140],
        [-3.0121, -0.0000, -3.2752, -1.1581]])
seqlogprob ln= False tensor([-9.9312, -7.4454])
seqlogprob ln= True tensor([-2.4828, -2.4818])
  sum tl*mask tensor([-9.9312, -7.4454])
  /masksum tensor([-2.4828, -2.4818])

Probe kl estimators

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(3)
def rn(*s): return torch.randn(*s, generator=g)
logp=rn(5); ref=rn(5)
d=logp-ref
for est in ['k1','k2','k3','mse','abs','kl']:
    try:
        print('kl_penalty',est, R.kl_penalty(logp,ref,est))
    except Exception as e:
        print('kl_penalty',est,'ERR',repr(e)[:60])
print('k1=d', d)
print('k2=.5d^2', 0.5*d*d)
print('k3=exp(-d)-1+d', torch.exp(-d)-1+d)
print('reverse_kl', R.reverse_kl(logp,ref))
print('symmetric_kl', R.symmetric_kl(logp,ref))
"
output
kl_penalty k1 tensor([ 2.1715, -0.1627, -0.9221, 0.8215, -0.9312])
kl_penalty k2 tensor([2.3578, 0.0132, 0.4252, 0.3374, 0.4336])
kl_penalty k3 tensor([1.2855, 0.0140, 0.5925, 0.2612, 0.6064])
kl_penalty mse ERR ValueError('mse')
kl_penalty abs ERR ValueError('abs')
kl_penalty kl ERR ValueError('kl')
k1=d tensor([ 2.1715, -0.1627, -0.9221, 0.8215, -0.9312])
k2=.5d^2 tensor([2.3578, 0.0132, 0.4252, 0.3374, 0.4336])
k3=exp(-d)-1+d tensor([1.2855, 0.0140, 0.5925, 0.2612, 0.6064])
reverse_kl tensor([1.2855, 0.0140, 0.5925, 0.2612, 0.6064])
symmetric_kl tensor([3.4429, 0.0133, 0.4562, 0.3568, 0.4659])

[stdout]
kl_penalty k1 tensor([ 2.1715, -0.1627, -0.9221, 0.8215, -0.9312])
kl_penalty k2 tensor([2.3578, 0.0132, 0.4252, 0.3374, 0.4336])
kl_penalty k3 tensor([1.2855, 0.0140, 0.5925, 0.2612, 0.6064])
kl_penalty mse ERR ValueError('mse')
kl_penalty abs ERR ValueError('abs')
kl_penalty kl ERR ValueError('kl')
k1=d tensor([ 2.1715, -0.1627, -0.9221, 0.8215, -0.9312])
k2=.5d^2 tensor([2.3578, 0.0132, 0.4252, 0.3374, 0.4336])
k3=exp(-d)-1+d tensor([1.2855, 0.0140, 0.5925, 0.2612, 0.6064])
reverse_kl tensor([1.2855, 0.0140, 0.5925, 0.2612, 0.6064])
symmetric_kl tensor([3.4429, 0.0133, 0.4562, 0.3568, 0.4659])

Probe dpo, ipo, bradley_terry

shell
python3 -c "
import reforacle as R, torch
import torch.nn.functional as F
g=torch.Generator().manual_seed(4)
def rn(*s): return torch.randn(*s, generator=g)
pc,pr,rc,rr=rn(6),rn(6),rn(6),rn(6)
beta=0.1; ls=0.05
print('dpo', R.dpo_loss(pc,pr,rc,rr,beta,ls))
logits=beta*((pc-rc)-(pr-rr))
loss=-(1-ls)*F.logsigmoid(logits)-ls*F.logsigmoid(-logits)
print(' mine mean', loss.mean())
print('ipo', R.ipo_loss(pc,pr,rc,rr,beta))
l2=((pc-rc)-(pr-rr)-1/(2*beta))**2
print(' mine', l2.mean())
cr,rj=rn(6),rn(6)
print('bt', R.bradley_terry_logit(cr,rj,beta))
print(' -logsig', (-F.logsigmoid(beta*(cr-rj))).mean(), 'nomean', -F.logsigmoid(beta*(cr-rj)))
print(' beta*(cr-rj)', beta*(cr-rj))
"
output
dpo tensor(0.6859)
 mine mean tensor(0.6859)
ipo tensor(25.4948)
 mine tensor(25.4948)
bt tensor([-0.1039, -0.1383, 0.1153, -0.0426, -0.0130, -0.0609])
 -logsig tensor(0.7145) nomean tensor([0.7464, 0.7647, 0.6371, 0.7147, 0.6997, 0.7241])
 beta*(cr-rj) tensor([-0.1039, -0.1383, 0.1153, -0.0426, -0.0130, -0.0609])

[stdout]
dpo tensor(0.6859)
 mine mean tensor(0.6859)
ipo tensor(25.4948)
 mine tensor(25.4948)
bt tensor([-0.1039, -0.1383, 0.1153, -0.0426, -0.0130, -0.0609])
 -logsig tensor(0.7145) nomean tensor([0.7464, 0.7647, 0.6371, 0.7147, 0.6997, 0.7241])
 beta*(cr-rj) tensor([-0.1039, -0.1383, 0.1153, -0.0426, -0.0130, -0.0609])

Probe importance_ratio, clip_fraction, normalize

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(5)
def rn(*s): return torch.randn(*s, generator=g)
logp,old=rn(6),rn(6)
print('importance_ratio clip=0.2', R.importance_ratio(logp,old,0.2))
r=torch.exp(logp-old)
print(' exp', r)
print(' clamp', r.clamp(1-0.2,1+0.2))
print('clip_fraction', R.clip_fraction(logp,old,0.2))
print(' |r-1|>0.2 mean', ((r-1).abs()>0.2).float().mean())
x=rn(7)
print('normalize eps=1e-8', R.normalize(x,1e-8))
print(' (x-mean)/(std+eps) unbiased', (x-x.mean())/(x.std()+1e-8))
print(' biased', (x-x.mean())/(x.std(unbiased=False)+1e-8))
"
output
importance_ratio clip=0.2 tensor([1.2000, 1.1720, 0.8000, 1.2000, 1.2000, 1.2000])
 exp tensor([2.4966, 1.1720, 0.3044, 4.7906, 1.4339, 3.4976])
 clamp tensor([1.2000, 1.1720, 0.8000, 1.2000, 1.2000, 1.2000])
clip_fraction tensor(0.8333)
 |r-1|>0.2 mean tensor(0.8333)
normalize eps=1e-8 tensor([ 0.1218, -1.3931, 0.9097, 0.7211, -1.2511, 1.0517, -0.1601])
 (x-mean)/(std+eps) unbiased tensor([ 0.1218, -1.3931, 0.9097, 0.7211, -1.2511, 1.0517, -0.1601])
 biased tensor([ 0.1316, -1.5048, 0.9826, 0.7789, -1.3513, 1.1360, -0.1729])

[stdout]
importance_ratio clip=0.2 tensor([1.2000, 1.1720, 0.8000, 1.2000, 1.2000, 1.2000])
 exp tensor([2.4966, 1.1720, 0.3044, 4.7906, 1.4339, 3.4976])
 clamp tensor([1.2000, 1.1720, 0.8000, 1.2000, 1.2000, 1.2000])
clip_fraction tensor(0.8333)
 |r-1|>0.2 mean tensor(0.8333)
normalize eps=1e-8 tensor([ 0.1218, -1.3931, 0.9097, 0.7211, -1.2511, 1.0517, -0.1601])
 (x-mean)/(std+eps) unbiased tensor([ 0.1218, -1.3931, 0.9097, 0.7211, -1.2511, 1.0517, -0.1601])
 biased tensor([ 0.1316, -1.5048, 0.9826, 0.7789, -1.3513, 1.1360, -0.1729])

Probe whiten variants

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(6)
def rn(*s): return torch.randn(*s, generator=g)
x=rn(8); mask=(torch.rand(8,generator=g)>0.3).float()
for sm in [False,True]:
    print('whiten sm=',sm, R.whiten(x,mask,sm))
    print('masked_whiten sm=',sm, R.masked_whiten(x,mask,sm))
# hypotheses
m=x.mean(); v=x.var(unbiased=False)
print('plain (x-m)/sqrt(v+1e-8)+m', (x-m)/torch.sqrt(v+1e-8)+m)
print('plain unbiased', (x-m)/(x.std())+m)
mm=(x*mask).sum()/mask.sum()
mv=((x-mm)**2*mask).sum()/mask.sum()
print('masked', (x-mm)/torch.sqrt(mv+1e-8))
print('masked+mean', (x-mm)/torch.sqrt(mv+1e-8)+mm)
"
output
whiten sm= False tensor([-1.5278, -0.8109, 0.5829, -0.5706, -1.1999, 0.1427, 1.1830, 0.5192])
masked_whiten sm= False tensor([-1.5278, -0.8109, 0.5829, -0.0107, -0.0107, -0.0107, 1.1830, 0.5192])
whiten sm= True tensor([-1.5171, -0.8002, 0.5936, -0.5598, -1.1892, 0.1534, 1.1937, 0.5300])
masked_whiten sm= True tensor([-1.5171, -0.8002, 0.5936, -0.0000, -0.0000, 0.0000, 1.1937, 0.5300])
plain (x-m)/sqrt(v+1e-8)+m tensor([-1.7252, -0.9256, 0.6286, -0.6576, -1.3595, 0.1377, 1.2979, 0.5577])
plain unbiased tensor([-1.6303, -0.8824, 0.5715, -0.6317, -1.2882, 0.1123, 1.1975, 0.5052])
masked tensor([-1.5171, -0.8002, 0.5936, -0.5598, -1.1892, 0.1534, 1.1937, 0.5300])
masked+mean tensor([-1.5278, -0.8109, 0.5829, -0.5706, -1.1999, 0.1427, 1.1830, 0.5192])

[stdout]
whiten sm= False tensor([-1.5278, -0.8109, 0.5829, -0.5706, -1.1999, 0.1427, 1.1830, 0.5192])
masked_whiten sm= False tensor([-1.5278, -0.8109, 0.5829, -0.0107, -0.0107, -0.0107, 1.1830, 0.5192])
whiten sm= True tensor([-1.5171, -0.8002, 0.5936, -0.5598, -1.1892, 0.1534, 1.1937, 0.5300])
masked_whiten sm= True tensor([-1.5171, -0.8002, 0.5936, -0.0000, -0.0000, 0.0000, 1.1937, 0.5300])
plain (x-m)/sqrt(v+1e-8)+m tensor([-1.7252, -0.9256, 0.6286, -0.6576, -1.3595, 0.1377, 1.2979, 0.5577])
plain unbiased tensor([-1.6303, -0.8824, 0.5715, -0.6317, -1.2882, 0.1123, 1.1975, 0.5052])
masked tensor([-1.5171, -0.8002, 0.5936, -0.5598, -1.1892, 0.1534, 1.1937, 0.5300])
masked+mean tensor([-1.5278, -0.8109, 0.5829, -0.5706, -1.1999, 0.1427, 1.1830, 0.5192])

Verify whiten var denom and advantage_mean_std

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(7)
def rn(*s): return torch.randn(*s, generator=g)
# verify whiten var denominator precisely
x=rn(8); mask=(torch.rand(8,generator=g)>0.3).float()
w=R.whiten(x,mask,True)
mm=(x*mask).sum()/mask.sum()
for name,den in [('N',mask.sum()),('N-1',mask.sum()-1)]:
    mv=((x-mm)**2*mask).sum()/den
    print('whiten den',name,'maxerr',(w-(x-mm)*torch.rsqrt(mv+1e-8)).abs().max().item())
mw=R.masked_whiten(x,mask,True)
mv=((x-mm)**2*mask).sum()/mask.sum()
print('masked_whiten==whitened*mask err',(mw-(x-mm)*torch.rsqrt(mv+1e-8)*mask).abs().max().item())
print('advantage_mean_std', R.advantage_mean_std(x,mask))
print(' mm,mstd_biased', mm, torch.sqrt(((x-mm)**2*mask).sum()/mask.sum()))
"
output
whiten den N maxerr 1.1920928955078125e-07
whiten den N-1 maxerr 0.14281439781188965
masked_whiten==whitened*mask err 1.1920928955078125e-07
advantage_mean_std tensor([0.4121, 0.9313])
 mm,mstd_biased tensor(0.4121) tensor(0.9313)

[stdout]
whiten den N maxerr 1.1920928955078125e-07
whiten den N-1 maxerr 0.14281439781188965
masked_whiten==whitened*mask err 1.1920928955078125e-07
advantage_mean_std tensor([0.4121, 0.9313])
 mm,mstd_biased tensor(0.4121) tensor(0.9313)

Probe gae, lambda_returns, discounted_returns

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(8)
def rn(*s): return torch.randn(*s, generator=g)
rewards=rn(5); values=rn(5); nv=torch.tensor(0.5); gamma=0.9; lam=0.95
print('gae', R.gae(rewards,values,nv,gamma,lam))
# compute gae
T=5; adv=torch.zeros(T); last=0.0
vals=torch.cat([values,nv.view(1)])
for t in reversed(range(T)):
    delta=rewards[t]+gamma*vals[t+1]-vals[t]
    last=delta+gamma*lam*last
    adv[t]=last
print(' mine',adv)
print('lambda_returns', R.lambda_returns(rewards,values,nv,gamma,lam))
print(' adv+values', adv+values)
print('discounted_returns', R.discounted_returns(rewards,gamma))
dr=torch.zeros(T); run=0.0
for t in reversed(range(T)):
    run=rewards[t]+gamma*run; dr[t]=run
print(' mine', dr)
"
output
gae tensor([ 0.3282, -0.5554, 0.7814, 1.8158, -0.8212])
 mine tensor([ 0.3282, -0.5554, 0.7814, 1.8158, -0.8212])
lambda_returns tensor([-0.3326, -0.6987, 0.6771, 0.2846, -0.1894])
 adv+values tensor([-0.3326, -0.6987, 0.6771, 0.2846, -0.1894])
discounted_returns tensor([-0.5819, -0.9480, 0.3610, -0.1574, -0.6394])
 mine tensor([-0.5819, -0.9480, 0.3610, -0.1574, -0.6394])

[stdout]
gae tensor([ 0.3282, -0.5554, 0.7814, 1.8158, -0.8212])
 mine tensor([ 0.3282, -0.5554, 0.7814, 1.8158, -0.8212])
lambda_returns tensor([-0.3326, -0.6987, 0.6771, 0.2846, -0.1894])
 adv+values tensor([-0.3326, -0.6987, 0.6771, 0.2846, -0.1894])
discounted_returns tensor([-0.5819, -0.9480, 0.3610, -0.1574, -0.6394])
 mine tensor([-0.5819, -0.9480, 0.3610, -0.1574, -0.6394])

Probe grpo/rloo advantages, group_mean_baseline

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(9)
def rn(*s): return torch.randn(*s, generator=g)
rewards=rn(6); gs=3
for sbs in [False,True]:
    print('grpo_adv sbs=',sbs, R.grpo_advantages(rewards,gs,sbs))
r=rewards.view(-1,gs)
mean=r.mean(1,keepdim=True)
std=r.std(1,keepdim=True)  # unbiased
stdb=r.std(1,unbiased=False,keepdim=True)
print(' centered', (r-mean).view(-1))
print(' /std_unb', ((r-mean)/(std+1e-4)).view(-1))
print(' /std_bias', ((r-mean)/(stdb+1e-4)).view(-1))
print('rloo_adv', R.rloo_advantages(rewards,gs))
n=gs
print(' loo', ((r*n-r.sum(1,keepdim=True))/(n-1)*0 + (r - (r.sum(1,keepdim=True)-r)/(n-1))).view(-1))
print('group_mean_baseline', R.group_mean_baseline(rewards,gs))
print(' mean broadcast', mean.expand(-1,gs).reshape(-1))
"
output
grpo_adv sbs= False tensor([-0.5303, 1.3363, -0.8060, -0.3319, 0.6402, -0.3082])
grpo_adv sbs= True tensor([-0.5572, 1.4043, -0.8470, -0.7331, 1.4139, -0.6808])
 centered tensor([-0.5303, 1.3363, -0.8060, -0.3319, 0.6402, -0.3082])
 /std_unb tensor([-0.4549, 1.1465, -0.6915, -0.5985, 1.1542, -0.5557])
 /std_bias tensor([-0.5572, 1.4041, -0.8469, -0.7330, 1.4136, -0.6806])
rloo_adv tensor([-0.7954, 2.0044, -1.2090, -0.4979, 0.9602, -0.4623])
 loo tensor([-0.7954, 2.0044, -1.2090, -0.4979, 0.9602, -0.4623])
group_mean_baseline tensor([-0.5303, 1.3363, -0.8060, -0.3319, 0.6402, -0.3082])
 mean broadcast tensor([0.5750, 0.5750, 0.5750, 0.6779, 0.6779, 0.6779])

[stdout]
grpo_adv sbs= False tensor([-0.5303, 1.3363, -0.8060, -0.3319, 0.6402, -0.3082])
grpo_adv sbs= True tensor([-0.5572, 1.4043, -0.8470, -0.7331, 1.4139, -0.6808])
 centered tensor([-0.5303, 1.3363, -0.8060, -0.3319, 0.6402, -0.3082])
 /std_unb tensor([-0.4549, 1.1465, -0.6915, -0.5985, 1.1542, -0.5557])
 /std_bias tensor([-0.5572, 1.4041, -0.8469, -0.7330, 1.4136, -0.6806])
rloo_adv tensor([-0.7954, 2.0044, -1.2090, -0.4979, 0.9602, -0.4623])
 loo tensor([-0.7954, 2.0044, -1.2090, -0.4979, 0.9602, -0.4623])
group_mean_baseline tensor([-0.5303, 1.3363, -0.8060, -0.3319, 0.6402, -0.3082])
 mean broadcast tensor([0.5750, 0.5750, 0.5750, 0.6779, 0.6779, 0.6779])

Nail grpo std eps and biased/unbiased

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(9)
def rn(*s): return torch.randn(*s, generator=g)
rewards=rn(6); gs=3
ref=R.grpo_advantages(rewards,gs,True)
r=rewards.view(-1,gs); mean=r.mean(1,keepdim=True)
stdb=r.std(1,unbiased=False,keepdim=True)
stdu=r.std(1,keepdim=True)
for eps in [1e-4,1e-5,1e-6,1e-8]:
    print('bias eps',eps,(ref-((r-mean)/(stdb+eps)).view(-1)).abs().max().item())
    print('unb  eps',eps,(ref-((r-mean)/(stdu+eps)).view(-1)).abs().max().item())
"
output
bias eps 0.0001 0.000308990478515625
unb  eps 0.0001 0.25965869426727295
bias eps 1e-05 2.8133392333984375e-05
unb  eps 1e-05 0.2594712972640991
bias eps 1e-06 0.0
unb  eps 1e-06 0.25945258140563965
bias eps 1e-08 3.0994415283203125e-06
unb  eps 1e-08 0.2594505548477173

[stdout]
bias eps 0.0001 0.000308990478515625
unb  eps 0.0001 0.25965869426727295
bias eps 1e-05 2.8133392333984375e-05
unb  eps 1e-05 0.2594712972640991
bias eps 1e-06 0.0
unb  eps 1e-06 0.25945258140563965
bias eps 1e-08 3.0994415283203125e-06
unb  eps 1e-08 0.2594505548477173

Probe value_loss, huber, clipped_pg

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(10)
def rn(*s): return torch.randn(*s, generator=g)
values=rn(6); old=rn(6); returns=rn(6); clip=0.2
print('value_loss', R.value_loss(values,old,returns,clip))
vc=old+(values-old).clamp(-clip,clip)
l=0.5*torch.max((values-returns)**2,(vc-returns)**2)
print(' mean', l.mean(), 'sum', l.sum())
print('huber d=1', R.huber_value_loss(values,returns,1.0))
import torch.nn.functional as F
print(' F.huber mean', F.huber_loss(values,returns,delta=1.0))
print(' smooth_l1', F.smooth_l1_loss(values,returns))
logp,old2,adv=rn(6),rn(6),rn(6); mask=(torch.rand(6,generator=g)>0.3).float()
print('clipped_pg', R.clipped_pg_loss(logp,old2,adv,mask,0.2,0.3))
r=torch.exp(logp-old2)
l1=-adv*r; l2=-adv*r.clamp(1-0.2,1+0.3)
pg=torch.max(l1,l2)
print(' masked_mean', (pg*mask).sum()/mask.sum())
"
output
value_loss tensor(2.3118)
 mean tensor(2.3118) sum tensor(13.8707)
huber d=1 tensor(1.2284)
 F.huber mean tensor(1.2284)
 smooth_l1 tensor(1.2284)
clipped_pg tensor(-0.2012)
 masked_mean tensor(-0.2012)

[stdout]
value_loss tensor(2.3118)
 mean tensor(2.3118) sum tensor(13.8707)
huber d=1 tensor(1.2284)
 F.huber mean tensor(1.2284)
 smooth_l1 tensor(1.2284)
clipped_pg tensor(-0.2012)
 masked_mean tensor(-0.2012)

Probe ce, smoothed_nll, logprob_at_temp, argmax, mode, median

shell
python3 -c "
import reforacle as R, torch
import torch.nn.functional as F
g=torch.Generator().manual_seed(11)
def rn(*s): return torch.randn(*s, generator=g)
logits=rn(3,7); labels=torch.randint(0,7,(3,),generator=g)
print('cross_entropy', R.cross_entropy(rn(4,7),torch.randint(0,7,(4,),generator=g),-100))
li=rn(4,7); lb=torch.randint(0,7,(4,),generator=g); lb[1]=-100
print('ce ignore', R.cross_entropy(li,lb,-100), F.cross_entropy(li,lb,ignore_index=-100))
print('smoothed_nll s=0.1', R.smoothed_nll(li,lb,0.1))
print(' F.ce ls', F.cross_entropy(li,lb,ignore_index=-100,label_smoothing=0.1))
print('logprob_at_temp t=2', R.logprob_at_temperature(rn(2,5),torch.randint(0,5,(2,),generator=g),2.0))
print('argmax_tokens', R.argmax_tokens(logits), logits.argmax(-1))
print('mode_label', R.mode_label(torch.tensor([1,2,2,3,3,3,1])))
print('median_reward', R.median_reward(torch.tensor([1.,3.,2.,4.])), torch.median(torch.tensor([1.,3.,2.,4.])))
"
output
Exit code 1
cross_entropy tensor(2.2595)
ce ignore tensor(1.3769) tensor(1.3769)
Traceback (most recent call last):
  File "<string>", line 10, in <module>
  File "reforacle.py", line 206, in reforacle.smoothed_nll
RuntimeError: index -100 is out of bounds for dimension 1 with size 7

[error] tool reported failure

Probe smoothed_nll formula

shell
python3 -c "
import reforacle as R, torch
import torch.nn.functional as F
g=torch.Generator().manual_seed(12)
def rn(*s): return torch.randn(*s, generator=g)
li=rn(4,7); lb=torch.randint(0,7,(4,),generator=g)
for s in [0.0,0.1]:
    print('smoothed_nll s=',s, R.smoothed_nll(li,lb,s))
    lp=torch.log_softmax(li,-1)
    nll=-lp.gather(-1,lb.unsqueeze(-1)).squeeze(-1)
    smooth=-lp.mean(-1)
    print(' (1-s)nll+s*smooth mean', ((1-s)*nll+s*smooth).mean())
    print(' F.ce ls', F.cross_entropy(li,lb,label_smoothing=s))
print('logprob_at_temp t=2', R.logprob_at_temperature(rn(2,5),torch.tensor([1,3]),2.0))
lg=rn(2,5)  # note regenerated; recompute
"
output
smoothed_nll s= 0.0 tensor(1.9227)
 (1-s)nll+s*smooth mean tensor(1.9227)
 F.ce ls tensor(1.9227)
smoothed_nll s= 0.1 tensor(1.9500)
 (1-s)nll+s*smooth mean tensor(1.9500)
 F.ce ls tensor(1.9500)
logprob_at_temp t=2 tensor([-2.0789, -2.0317])

[stdout]
smoothed_nll s= 0.0 tensor(1.9227)
 (1-s)nll+s*smooth mean tensor(1.9227)
 F.ce ls tensor(1.9227)
smoothed_nll s= 0.1 tensor(1.9500)
 (1-s)nll+s*smooth mean tensor(1.9500)
 F.ce ls tensor(1.9500)
logprob_at_temp t=2 tensor([-2.0789, -2.0317])

Probe remaining utilities

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(13)
def rn(*s): return torch.randn(*s, generator=g)
print('mode_label', R.mode_label(torch.tensor([1,2,2,3,3,3,1])))
print('median_reward even', R.median_reward(torch.tensor([1.,3.,2.,4.])))
print('  torch.median', torch.median(torch.tensor([1.,3.,2.,4.])), 'quantile.5', torch.quantile(torch.tensor([1.,3.,2.,4.]),0.5))
v=rn(6); m=(torch.rand(6,generator=g)>0.4).float()
print('masked_max', R.masked_max(v,m)); print(' v',v,'m',m)
x=rn(10)
print('quantile_lower q=0.3', R.quantile_lower(x,0.3))
print(' torch.quantile lower', torch.quantile(x,0.3,interpolation='lower'), 'linear', torch.quantile(x,0.3))
print('cumulative_max', R.cumulative_max(torch.tensor([1.,3.,2.,5.,4.])))
print('bucketize', R.bucketize_reward(torch.tensor([0.1,0.5,0.9,1.5]),torch.tensor([0.2,0.6,1.0])))
print('pad_mask', R.pad_mask_from_lengths(torch.tensor([2,3,0]),4))
print('first_nonzero', R.first_nonzero_index(torch.tensor([[0.,0.,1.,0.],[1.,0.,0.,0.],[0.,0.,0.,0.]])))
"
output
mode_label tensor(3)
median_reward even tensor(2.5000)
  torch.median tensor(2.) quantile.5 tensor(2.5000)
masked_max tensor(1.4103)
 v tensor([-0.1117, -1.1456, 1.4103, 1.3654, 1.1206, -0.9622]) m tensor([1., 1., 1., 0., 0., 0.])
quantile_lower q=0.3 tensor(-0.6628)
 torch.quantile lower tensor(-0.6628) linear tensor(-0.6080)
cumulative_max tensor([1., 1., 3., 3., 5.])
bucketize tensor([0, 1, 2, 3])
pad_mask tensor([[1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [0., 0., 0., 0.]])
first_nonzero tensor([ 2, 0, -1])

[stdout]
mode_label tensor(3)
median_reward even tensor(2.5000)
  torch.median tensor(2.) quantile.5 tensor(2.5000)
masked_max tensor(1.4103)
 v tensor([-0.1117, -1.1456, 1.4103, 1.3654, 1.1206, -0.9622]) m tensor([1., 1., 1., 0., 0., 0.])
quantile_lower q=0.3 tensor(-0.6628)
 torch.quantile lower tensor(-0.6628) linear tensor(-0.6080)
cumulative_max tensor([1., 1., 3., 3., 5.])
bucketize tensor([0, 1, 2, 3])
pad_mask tensor([[1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [0., 0., 0., 0.]])
first_nonzero tensor([ 2, 0, -1])

Probe top_k, top_p, 2D behaviors

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(14)
def rn(*s): return torch.randn(*s, generator=g)
# top_k_mask
logits=rn(2,6)
print('logits',logits)
print('top_k_mask k=3', R.top_k_mask(logits,3))
# top_p_mask
probs=torch.softmax(rn(2,6),-1)
print('probs',probs)
print('top_p_mask p=0.7', R.top_p_mask(probs,0.7))
# 2D masked_max
v=rn(2,4); m=(torch.rand(2,4,generator=g)>0.4).float()
print('masked_max 2D', R.masked_max(v,m))
print('cumulative_max 2D', R.cumulative_max(rn(2,4)))
print('first_nonzero 1D', R.first_nonzero_index(torch.tensor([0.,0.,1.,0.])))
print('median odd', R.median_reward(torch.tensor([3.,1.,2.])))
"
output
logits tensor([[-1.0141, -0.3720, -0.7516, -0.8623, -0.3270, 0.5212],
        [ 1.2622, -1.4680, -0.1037, 0.5177, -1.0845, -2.0901]])
top_k_mask k=3 tensor([[False, True, False, False, True, True],
        [ True, False, True, True, False, False]])
probs tensor([[0.1499, 0.1438, 0.1177, 0.1226, 0.3846, 0.0813],
        [0.7357, 0.0484, 0.1154, 0.0282, 0.0283, 0.0440]])
top_p_mask p=0.7 tensor([[0.1872, 0.1795, 0.0000, 0.1531, 0.4802, 0.0000],
        [1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]])
masked_max 2D tensor(1.1964)
cumulative_max 2D tensor([[-0.7920, 0.4916, 2.0040, -0.2893],
        [-0.7920, 0.4916, 2.0040, -0.2893]])
first_nonzero 1D tensor(2)
median odd tensor(2.)

[stdout]
logits tensor([[-1.0141, -0.3720, -0.7516, -0.8623, -0.3270, 0.5212],
        [ 1.2622, -1.4680, -0.1037, 0.5177, -1.0845, -2.0901]])
top_k_mask k=3 tensor([[False, True, False, False, True, True],
        [ True, False, True, True, False, False]])
probs tensor([[0.1499, 0.1438, 0.1177, 0.1226, 0.3846, 0.0813],
        [0.7357, 0.0484, 0.1154, 0.0282, 0.0283, 0.0440]])
top_p_mask p=0.7 tensor([[0.1872, 0.1795, 0.0000, 0.1531, 0.4802, 0.0000],
        [1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]])
masked_max 2D tensor(1.1964)
cumulative_max 2D tensor([[-0.7920, 0.4916, 2.0040, -0.2893],
        [-0.7920, 0.4916, 2.0040, -0.2893]])
first_nonzero 1D tensor(2)
median odd tensor(2.)

Determine cummax dim, mode tie-break, masked_max edge

shell
python3 -c "
import reforacle as R, torch
x=torch.tensor([[1.,3.,2.,0.],[5.,1.,4.,2.]])
print('cummax input',x)
print('reforacle', R.cumulative_max(x))
print('dim0', torch.cummax(x,0).values)
print('dim1', torch.cummax(x,-1).values)
print('flat', torch.cummax(x.flatten(),0).values.view_as(x))
# mode tie
print('mode tie [1,1,2,2]', R.mode_label(torch.tensor([1,1,2,2])))
print('mode tie [2,2,1,1]', R.mode_label(torch.tensor([2,2,1,1])))
print('torch.mode', torch.mode(torch.tensor([1,1,2,2])).values)
# masked_max mask all zero
print('masked_max allzero', R.masked_max(torch.tensor([1.,2.]),torch.tensor([0.,0.])))
"
output
cummax input tensor([[1., 3., 2., 0.],
        [5., 1., 4., 2.]])
reforacle tensor([[1., 3., 2., 0.],
        [1., 3., 2., 0.]])
dim0 tensor([[1., 3., 2., 0.],
        [5., 3., 4., 2.]])
dim1 tensor([[1., 3., 3., 3.],
        [5., 5., 5., 5.]])
flat tensor([[1., 3., 3., 3.],
        [5., 5., 5., 5.]])
mode tie [1,1,2,2] tensor(2)
mode tie [2,2,1,1] tensor(2)
torch.mode tensor(1)
masked_max allzero tensor(0.)

[stdout]
cummax input tensor([[1., 3., 2., 0.],
        [5., 1., 4., 2.]])
reforacle tensor([[1., 3., 2., 0.],
        [1., 3., 2., 0.]])
dim0 tensor([[1., 3., 2., 0.],
        [5., 3., 4., 2.]])
dim1 tensor([[1., 3., 3., 3.],
        [5., 5., 5., 5.]])
flat tensor([[1., 3., 3., 3.],
        [5., 5., 5., 5.]])
mode tie [1,1,2,2] tensor(2)
mode tie [2,2,1,1] tensor(2)
torch.mode tensor(1)
masked_max allzero tensor(0.)

Nail cumulative_max semantics

shell
python3 -c "
import reforacle as R, torch
print('1D', R.cumulative_max(torch.tensor([1.,3.,2.,5.,4.])))
x3=torch.tensor([[1.,9.],[5.,2.],[3.,8.]])
print('3x2 input',x3.tolist())
print('reforacle', R.cumulative_max(x3))
x=torch.tensor([[1.,3.,2.,0.],[5.,1.,4.,2.]])
print('cummax dim0 col-run?', torch.cummax(x,0).values)
# maybe it does cummax over flattened along axis then reshape column-major?
"
output
1D tensor([1., 1., 3., 3., 5.])
3x2 input [[1.0, 9.0], [5.0, 2.0], [3.0, 8.0]]
reforacle tensor([[1., 9.],
        [1., 9.],
        [5., 9.]])
cummax dim0 col-run? tensor([[1., 3., 2., 0.],
        [5., 3., 4., 2.]])

[stdout]
1D tensor([1., 1., 3., 3., 5.])
3x2 input [[1.0, 9.0], [5.0, 2.0], [3.0, 8.0]]
reforacle tensor([[1., 9.],
        [1., 9.],
        [5., 9.]])
cummax dim0 col-run? tensor([[1., 3., 2., 0.],
        [5., 3., 4., 2.]])

Probe mode tie-break and masked_max negatives

shell
python3 -c "
import reforacle as R, torch
print('mode [3,3,1,1,2]', R.mode_label(torch.tensor([3,3,1,1,2])))
print('mode [5,1,1]', R.mode_label(torch.tensor([5,1,1])))
print('mode [1,1,5,5,2,2]', R.mode_label(torch.tensor([1,1,5,5,2,2])))
# masked_max negative unmasked-large
print('mm neg', R.masked_max(torch.tensor([-1.,-2.,5.]),torch.tensor([1.,1.,0.])))
print('mm neg2', R.masked_max(torch.tensor([-3.,-2.]),torch.tensor([1.,1.])))
"
output
mode [3,3,1,1,2] tensor(3)
mode [5,1,1] tensor(1)
mode [1,1,5,5,2,2] tensor(5)
mm neg tensor(-1.)
mm neg2 tensor(-2.)

[stdout]
mode [3,3,1,1,2] tensor(3)
mode [5,1,1] tensor(1)
mode [1,1,5,5,2,2] tensor(5)
mm neg tensor(-1.)
mm neg2 tensor(-2.)

Probe dpo_sequence_loss

shell
python3 -c "
import reforacle as R, torch
import torch.nn.functional as F
g=torch.Generator().manual_seed(20)
def rn(*s): return torch.randn(*s, generator=g)
B,T,V=2,4,5
pcl,prl,rcl,rrl=rn(B,T,V),rn(B,T,V),rn(B,T,V),rn(B,T,V)
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(pcl,prl,rcl,rrl,cl,rl,cm,rm,beta,ls)
print('dpo_seq', out)
def seqlp(logits,labels,mask):
    lp=torch.log_softmax(logits,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
    return (lp*mask).sum(-1)
pc=seqlp(pcl,cl,cm); pr=seqlp(prl,rl,rm); rc=seqlp(rcl,cl,cm); rr=seqlp(rrl,rl,rm)
logits=beta*((pc-rc)-(pr-rr))
loss=(-(1-ls)*F.logsigmoid(logits)-ls*F.logsigmoid(-logits)).mean()
print('mine', loss)
"
output
dpo_seq tensor(0.5428)
mine tensor(0.5428)

[stdout]
dpo_seq tensor(0.5428)
mine tensor(0.5428)

Probe rloo_objective

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(21)
def rn(*s): return torch.randn(*s, generator=g)
B,T,V=6,4,5; gs=3
logits=rn(B,T,V); old=rn(B,T,V)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
rewards=rn(B)
cl,ch=0.2,0.3
out=R.rloo_objective(logits,old,labels,mask,rewards,gs,cl,ch)
print('rloo_obj', out)
def tlp(lg): return torch.log_softmax(lg,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
lp=tlp(logits); olp=tlp(old)
r=rewards.view(-1,gs); n=gs
adv=(r - (r.sum(1,keepdim=True)-r)/(n-1)).view(-1)  # (B,)
advt=adv.unsqueeze(-1)  # broadcast to T
ratio=torch.exp(lp-olp)
l1=-advt*ratio; l2=-advt*ratio.clamp(1-cl,1+ch)
pg=torch.max(l1,l2)
print('masked_mean', (pg*mask).sum()/mask.sum())
"
output
rloo_obj tensor(0.2949)
masked_mean tensor(0.2949)

[stdout]
rloo_obj tensor(0.2949)
masked_mean tensor(0.2949)

Probe reverse_kl_objective

shell
python3 -c "
import reforacle as R, torch
import torch.nn.functional as F
g=torch.Generator().manual_seed(22)
def rn(*s): return torch.randn(*s, generator=g)
B,T,V=3,4,5
logits=rn(B,T,V); ref=rn(B,T,V)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=rn(B)
beta=0.1
out=R.reverse_kl_objective(logits,ref,labels,mask,adv,beta)
print('rkl_obj', out)
def tlp(lg): return torch.log_softmax(lg,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
lp=tlp(logits); rlp=tlp(ref)
d=lp-rlp
kl=torch.exp(-d)-1+d  # reverse_kl k3 per token
pg=-adv.unsqueeze(-1)*lp
print('A: mm(-adv*lp)+beta*mm(kl)', (pg*mask).sum()/mask.sum()+beta*((kl*mask).sum()/mask.sum()))
print('B: mm(-adv*lp+beta*kl)', ((pg+beta*kl)*mask).sum()/mask.sum())
print('adv broadcast? maybe per token adv shape', adv.shape)
"
output
rkl_obj tensor(1.0201)
A: mm(-adv*lp)+beta*mm(kl) tensor(1.0201)
B: mm(-adv*lp+beta*kl) tensor(1.0201)
adv broadcast? maybe per token adv shape torch.Size([3])

[stdout]
rkl_obj tensor(1.0201)
A: mm(-adv*lp)+beta*mm(kl) tensor(1.0201)
B: mm(-adv*lp+beta*kl) tensor(1.0201)
adv broadcast? maybe per token adv shape torch.Size([3])

Probe grpo_objective composition

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(23)
def rn(*s): return torch.randn(*s, generator=g)
B,T,V=6,4,5; gs=3
logits=rn(B,T,V); old=rn(B,T,V); ref=rn(B,T,V)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
rewards=rn(B)
beta,cl,ch=0.04,0.2,0.3
for sbs in [False,True]:
  for est in ['k1','k2','k3']:
    out=R.grpo_objective(logits,old,ref,labels,mask,rewards,gs,beta,cl,ch,sbs,est)
    def tlp(lg): return torch.log_softmax(lg,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
    lp=tlp(logits); olp=tlp(old); rlp=tlp(ref)
    r=rewards.view(-1,gs); m=r.mean(1,keepdim=True)
    if sbs: adv=((r-m)/(r.std(1,unbiased=False,keepdim=True)+1e-6)).view(-1)
    else: adv=(r-m).view(-1)
    advt=adv.unsqueeze(-1)
    ratio=torch.exp(lp-olp)
    pg=torch.max(-advt*ratio,-advt*ratio.clamp(1-cl,1+ch))
    pgl=(pg*mask).sum()/mask.sum()
    d=lp-rlp
    kl={'k1':d,'k2':0.5*d*d,'k3':torch.exp(-d)-1+d}[est]
    kll=(kl*mask).sum()/mask.sum()
    mine=pgl+beta*kll
    print(sbs,est,'ref',out.item(),'mine',mine.item(),'diff',(out-mine).abs().item())
"
output
False k1 ref 2.7708613872528076 mine 2.7708611488342285 diff 2.384185791015625e-07
False k2 ref 2.784942388534546 mine 2.784942150115967 diff 2.384185791015625e-07
False k3 ref 2.7873611450195312 mine 2.787360906600952 diff 2.384185791015625e-07
True k1 ref 2.799774408340454 mine 2.799774169921875 diff 2.384185791015625e-07
True k2 ref 2.8138551712036133 mine 2.8138551712036133 diff 0.0
True k3 ref 2.8162739276885986 mine 2.8162739276885986 diff 0.0

[stdout]
False k1 ref 2.7708613872528076 mine 2.7708611488342285 diff 2.384185791015625e-07
False k2 ref 2.784942388534546 mine 2.784942150115967 diff 2.384185791015625e-07
False k3 ref 2.7873611450195312 mine 2.787360906600952 diff 2.384185791015625e-07
True k1 ref 2.799774408340454 mine 2.799774169921875 diff 2.384185791015625e-07
True k2 ref 2.8138551712036133 mine 2.8138551712036133 diff 0.0
True k3 ref 2.8162739276885986 mine 2.8162739276885986 diff 0.0

Probe ppo_objective

shell
python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(24)
def rn(*s): return torch.randn(*s, generator=g)
T=6
rewards=rn(T); values=rn(T); old_values=rn(T); logp=rn(T); old_logp=rn(T)
nv=torch.tensor(0.3); gamma,lam=0.99,0.95; cl,ch=0.2,0.2; vfclip=0.2; vfcoef=0.5
out=R.ppo_objective(rewards,values,old_values,logp,old_logp,nv,gamma,lam,cl,ch,vfclip,vfcoef)
print('ppo', out)
# gae
adv=torch.zeros(T); last=0.0; vals=torch.cat([values,nv.view(1)])
for t in reversed(range(T)):
    delta=rewards[t]+gamma*vals[t+1]-vals[t]; last=delta+gamma*lam*last; adv[t]=last
returns=adv+values
ratio=torch.exp(logp-old_logp)
def pg(a):
    return torch.max(-a*ratio,-a*ratio.clamp(1-cl,1+ch)).mean()
vc=old_values+(values-old_values).clamp(-vfclip,vfclip)
vl=0.5*torch.max((values-returns)**2,(vc-returns)**2).mean()
# whiten adv?
advw=(adv-adv.mean())/(adv.std(unbiased=False)+1e-8)  # whiten shift_mean True
print('raw adv:', pg(adv)+vfcoef*vl)
print('whiten adv:', pg(advw)+vfcoef*vl)
"
output
ppo tensor(0.7886)
raw adv: tensor(0.6986)
whiten adv: tensor(0.7886)

[stdout]
ppo tensor(0.7886)
raw adv: tensor(0.6986)
whiten adv: tensor(0.7886)
write · /app/objectives.py
from __future__ import annotations

import math

import torch
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# internal helpers
# ---------------------------------------------------------------------------
def _token_logprobs(logits, labels):
    lp = torch.log_softmax(logits, dim=-1)
    return lp.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)


def _seq_logprob(logits, labels, mask):
    return (_token_logprobs(logits, labels) * mask).sum(dim=-1)


# ---------------------------------------------------------------------------
# reductions / masking
# ---------------------------------------------------------------------------
def masked_mean(values, mask, axis=None):
    if axis is None:
        num = (values * mask).sum()
        den = mask.sum()
    else:
        num = (values * mask).sum(dim=axis)
        den = mask.sum(dim=axis)
    return torch.where(den == 0, torch.zeros_like(num), num / den)


def masked_sum(values, mask, axis=None):
    if axis is None:
        return (values * mask).sum()
    return (values * mask).sum(dim=axis)


def masked_max(values, mask):
    filled = values.masked_fill(mask == 0, float("-inf"))
    m = filled.max()
    return torch.where(mask.sum() > 0, m, torch.zeros_like(m))


def logsumexp(x, axis):
    return torch.logsumexp(x, dim=axis)


def log_softmax(x, axis):
    return torch.log_softmax(x, dim=axis)


# ---------------------------------------------------------------------------
# logprobs / entropy
# ---------------------------------------------------------------------------
def token_logprobs(logits, labels):
    return _token_logprobs(logits, labels)


def selective_logprobs(logits, labels, mask):
    return (_token_logprobs(logits, labels) * mask).sum(dim=-1)


def sequence_logprob(logits, labels, mask, length_normalize):
    tl = _token_logprobs(logits, labels)
    if length_normalize:
        return masked_mean(tl, mask, axis=-1)
    return (tl * mask).sum(dim=-1)


def logprob_at_temperature(logits, labels, temperature):
    lp = torch.log_softmax(logits / temperature, dim=-1)
    return lp.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)


def entropy(logits, mask):
    lp = torch.log_softmax(logits, dim=-1)
    p = lp.exp()
    ent = -(p * lp).sum(dim=-1)
    return masked_mean(ent, mask)


def normalized_entropy(logits, mask):
    v = logits.shape[-1]
    return entropy(logits, mask) / math.log(v)


def cross_entropy(logits, labels, ignore_index):
    return F.cross_entropy(logits, labels.long(), ignore_index=ignore_index)


def smoothed_nll(logits, labels, smoothing):
    lp = torch.log_softmax(logits, dim=-1)
    nll = -lp.gather(-1, labels.long().unsqueeze(-1)).squeeze(-1)
    smooth = -lp.mean(dim=-1)
    return ((1 - smoothing) * nll + smoothing * smooth).mean()


# ---------------------------------------------------------------------------
# preference objectives
# ---------------------------------------------------------------------------
def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
    logits = beta * ((pc - rc) - (pr - rr))
    loss = -(1 - label_smoothing) * F.logsigmoid(logits) - label_smoothing * F.logsigmoid(-logits)
    return loss.mean()


def ipo_loss(pc, pr, rc, rr, beta):
    diff = (pc - rc) - (pr - rr)
    return ((diff - 1.0 / (2.0 * beta)) ** 2).mean()


def bradley_terry_logit(chosen_reward, rejected_reward, beta):
    return beta * (chosen_reward - rejected_reward)


def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
                      chosen_labels, rejected_labels, chosen_mask, rejected_mask,
                      beta, label_smoothing):
    pc = _seq_logprob(pc_logits, chosen_labels, chosen_mask)
    pr = _seq_logprob(pr_logits, rejected_labels, rejected_mask)
    rc = _seq_logprob(rc_logits, chosen_labels, chosen_mask)
    rr = _seq_logprob(rr_logits, rejected_labels, rejected_mask)
    return dpo_loss(pc, pr, rc, rr, beta, label_smoothing)


# ---------------------------------------------------------------------------
# advantages / baselines
# ---------------------------------------------------------------------------
def grpo_advantages(rewards, group_size, scale_by_std):
    r = rewards.view(-1, group_size)
    mean = r.mean(dim=1, keepdim=True)
    adv = r - mean
    if scale_by_std:
        std = r.std(dim=1, unbiased=False, keepdim=True)
        adv = adv / (std + 1e-6)
    return adv.reshape(rewards.shape)


def rloo_advantages(rewards, group_size):
    r = rewards.view(-1, group_size)
    n = group_size
    total = r.sum(dim=1, keepdim=True)
    adv = r - (total - r) / (n - 1)
    return adv.reshape(rewards.shape)


def group_mean_baseline(rewards, group_size):
    r = rewards.view(-1, group_size)
    mean = r.mean(dim=1, keepdim=True)
    return (r - mean).reshape(rewards.shape)


def advantage_mean_std(advantages, mask):
    mean = masked_mean(advantages, mask)
    var = masked_mean((advantages - mean) ** 2, mask)
    return torch.stack([mean, torch.sqrt(var)])


# ---------------------------------------------------------------------------
# returns
# ---------------------------------------------------------------------------
def gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[-1]
    adv = torch.zeros_like(rewards)
    last = torch.zeros_like(rewards[..., 0])
    for t in reversed(range(T)):
        v_next = values[..., t + 1] if t + 1 < T else next_value
        delta = rewards[..., t] + gamma * v_next - values[..., t]
        last = delta + gamma * lam * last
        adv[..., t] = last
    return adv


def lambda_returns(rewards, values, next_value, gamma, lam):
    return gae(rewards, values, next_value, gamma, lam) + values


def discounted_returns(rewards, gamma):
    T = rewards.shape[-1]
    out = torch.zeros_like(rewards)
    run = torch.zeros_like(rewards[..., 0])
    for t in reversed(range(T)):
        run = rewards[..., t] + gamma * run
        out[..., t] = run
    return out


# ---------------------------------------------------------------------------
# KL estimators
# ---------------------------------------------------------------------------
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) - 1 + d
    raise ValueError(estimator)


def reverse_kl(logp, ref_logp):
    d = logp - ref_logp
    return torch.exp(-d) - 1 + d


def symmetric_kl(logp, ref_logp):
    return torch.cosh(logp - ref_logp) - 1


# ---------------------------------------------------------------------------
# policy-gradient / value losses
# ---------------------------------------------------------------------------
def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
    ratio = torch.exp(logp - old_logp)
    l1 = -advantages * ratio
    l2 = -advantages * ratio.clamp(1 - clip_low, 1 + clip_high)
    return masked_mean(torch.max(l1, l2), mask)


def value_loss(values, old_values, returns, clip):
    vclip = old_values + (values - old_values).clamp(-clip, clip)
    return (0.5 * torch.max((values - returns) ** 2, (vclip - returns) ** 2)).mean()


def huber_value_loss(values, returns, delta):
    return F.huber_loss(values, returns, delta=delta)


def importance_ratio(logp, old_logp, clip):
    return torch.exp(logp - old_logp).clamp(1 - clip, 1 + clip)


def clip_fraction(logp, old_logp, clip):
    ratio = torch.exp(logp - old_logp)
    return ((ratio - 1).abs() > clip).float().mean()


# ---------------------------------------------------------------------------
# whitening / normalization
# ---------------------------------------------------------------------------
def whiten(values, mask, shift_mean):
    mean = masked_mean(values, mask)
    var = masked_mean((values - mean) ** 2, mask)
    w = (values - mean) * torch.rsqrt(var + 1e-8)
    if not shift_mean:
        w = w + mean
    return w


def masked_whiten(values, mask, shift_mean):
    mean = masked_mean(values, mask)
    var = masked_mean((values - mean) ** 2, mask)
    w = (values - mean) * torch.rsqrt(var + 1e-8) * mask
    if not shift_mean:
        w = w + mean
    return w


def normalize(x, eps):
    return (x - x.mean()) / (x.std() + eps)


# ---------------------------------------------------------------------------
# composite objectives
# ---------------------------------------------------------------------------
def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
                   rewards, group_size, beta, clip_low, clip_high, scale_by_std,
                   kl_estimator):
    lp = _token_logprobs(logits, labels)
    olp = _token_logprobs(old_logits, labels)
    rlp = _token_logprobs(ref_logits, labels)
    adv = grpo_advantages(rewards, group_size, scale_by_std).unsqueeze(-1)
    pg = clipped_pg_loss(lp, olp, adv, completion_mask, clip_low, clip_high)
    kl = kl_penalty(lp, rlp, kl_estimator)
    return pg + beta * masked_mean(kl, 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_w = whiten(adv, torch.ones_like(adv), True)
    pg = clipped_pg_loss(logp, old_logp, adv_w, torch.ones_like(logp), clip_low, clip_high)
    vl = value_loss(values, old_values, returns, vf_clip)
    return pg + vf_coef * vl


def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
    lp = _token_logprobs(logits, labels)
    olp = _token_logprobs(old_logits, labels)
    adv = rloo_advantages(rewards, group_size).unsqueeze(-1)
    return clipped_pg_loss(lp, olp, adv, mask, clip_low, clip_high)


def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
    lp = _token_logprobs(logits, labels)
    rlp = _token_logprobs(ref_logits, labels)
    kl = reverse_kl(lp, rlp)
    per = -advantages.unsqueeze(-1) * lp + beta * kl
    return masked_mean(per, mask)


# ---------------------------------------------------------------------------
# sampling / utility
# ---------------------------------------------------------------------------
def top_p_mask(probs, p):
    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    cumsum = sorted_probs.cumsum(dim=-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)
    return out / out.sum(dim=-1, keepdim=True)


def top_k_mask(logits, k):
    _, idx = torch.topk(logits, k, dim=-1)
    mask = torch.zeros_like(logits, dtype=torch.bool)
    mask.scatter_(-1, idx, True)
    return mask


def argmax_tokens(logits):
    return logits.argmax(dim=-1)


def mode_label(labels):
    vals, counts = torch.unique(labels, return_counts=True)
    maxc = counts.max()
    return vals[counts == maxc].max()


def median_reward(rewards):
    return torch.quantile(rewards.float(), 0.5)


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
    has = nz.any(dim=-1)
    idx = nz.to(torch.float32).argmax(dim=-1)
    return torch.where(has, idx, torch.full_like(idx, -1))


def cumulative_max(x):
    cm = torch.cummax(x, dim=0).values
    return torch.cat([x[:1], cm[:-1]], dim=0)


def bucketize_reward(rewards, boundaries):
    return torch.bucketize(rewards, boundaries)
write · /app/check.py
import torch
import reforacle as R
import objectives as O

torch.manual_seed(0)

def rn(*s, g=None):
    return torch.randn(*s)

def cmp(name, ref, mine):
    if isinstance(ref, tuple):
        ok = all(cmp1(a, b) for a, b in zip(ref, mine))
    else:
        ok = cmp1(ref, mine)
    if not ok:
        print(f"FAIL {name}\n  ref ={ref}\n  mine={mine}")
    return ok

def cmp1(ref, mine):
    ref = torch.as_tensor(ref)
    mine = torch.as_tensor(mine)
    if ref.shape != mine.shape:
        print("  shape mismatch", ref.shape, mine.shape)
        return False
    if ref.dtype == torch.bool or not ref.is_floating_point():
        return torch.equal(ref, mine.to(ref.dtype)) if ref.dtype==mine.dtype else torch.equal(ref.long(), mine.long())
    return torch.allclose(ref, mine, atol=1e-5, rtol=1e-4, equal_nan=True)

def run_case(name, fn_ref, fn_mine, args):
    try:
        ref = fn_ref(*args)
    except Exception as e:
        print(f"  ref raised {name}: {e}")
        return True
    mine = fn_mine(*args)
    return cmp(name, ref, mine)

def G(*shape):
    return torch.randn(*shape)

def M(*shape, p=0.4):
    return (torch.rand(*shape) > p).float()

def L(hi, *shape):
    return torch.randint(0, hi, shape)

results = {}
def T(name, args_list):
    ok = True
    for args in args_list:
        ok = run_case(name, getattr(R, name), getattr(O, name), args) and ok
    results[name] = ok

for _ in range(5):
    T("masked_mean", [(G(8), M(8)), (G(3,5), M(3,5), 1), (G(4), torch.zeros(4))])
    T("masked_sum", [(G(8), M(8)), (G(3,5), M(3,5), 1), (G(3,5), M(3,5), 0)])
    T("masked_max", [(G(8), M(8)), (G(3,4), M(3,4)), (G(5), torch.zeros(5))])
    T("logsumexp", [(G(3,5), 1), (G(3,5), 0), (G(4), 0)])
    T("log_softmax", [(G(3,5), 1), (G(2,3,4), -1)])
    T("token_logprobs", [(G(2,4,6), L(6,2,4)), (G(3,5), L(5,3))])
    T("selective_logprobs", [(G(2,4,6), L(6,2,4), M(2,4))])
    T("sequence_logprob", [(G(2,4,6), L(6,2,4), M(2,4), False), (G(2,4,6), L(6,2,4), M(2,4), True)])
    T("logprob_at_temperature", [(G(2,4,6), L(6,2,4), 2.0), (G(3,5), L(5,3), 0.7)])
    T("entropy", [(G(3,5), M(3)), (G(2,4,6), M(2,4))])
    T("normalized_entropy", [(G(3,5), M(3)), (G(2,4,6), M(2,4))])
    T("cross_entropy", [(G(4,7), L(7,4), -100), (G(5,6), L(6,5), -100)])
    T("smoothed_nll", [(G(4,7), L(7,4), 0.1), (G(5,6), L(6,5), 0.0)])
    T("dpo_loss", [(G(6),G(6),G(6),G(6),0.1,0.0), (G(5),G(5),G(5),G(5),0.5,0.1)])
    T("ipo_loss", [(G(6),G(6),G(6),G(6),0.1)])
    T("bradley_terry_logit", [(G(6),G(6),0.1), (G(3),G(3),1.0)])
    T("dpo_sequence_loss", [(G(2,4,5),G(2,4,5),G(2,4,5),G(2,4,5),L(5,2,4),L(5,2,4),M(2,4),M(2,4),0.1,0.0)])
    T("grpo_advantages", [(G(6),3,False),(G(6),3,True),(G(9),3,True)])
    T("rloo_advantages", [(G(6),3),(G(9),3)])
    T("group_mean_baseline", [(G(6),3),(G(8),4)])
    T("advantage_mean_std", [(G(8),M(8)),(G(3,4),M(3,4))])
    T("gae", [(G(5),G(5),torch.tensor(0.5),0.9,0.95)])
    T("lambda_returns", [(G(5),G(5),torch.tensor(0.3),0.99,0.95)])
    T("discounted_returns", [(G(5),0.9),(G(7),0.95)])
    T("kl_penalty", [(G(5),G(5),'k1'),(G(5),G(5),'k2'),(G(5),G(5),'k3')])
    T("reverse_kl", [(G(5),G(5))])
    T("symmetric_kl", [(G(5),G(5))])
    T("clipped_pg_loss", [(G(6),G(6),G(6),M(6),0.2,0.3)])
    T("value_loss", [(G(6),G(6),G(6),0.2)])
    T("huber_value_loss", [(G(6),G(6),1.0),(G(6),G(6),0.5)])
    T("importance_ratio", [(G(6),G(6),0.2)])
    T("clip_fraction", [(G(6),G(6),0.2)])
    T("whiten", [(G(8),M(8),False),(G(8),M(8),True)])
    T("masked_whiten", [(G(8),M(8),False),(G(8),M(8),True)])
    T("normalize", [(G(7),1e-8)])
    T("grpo_objective", [(G(6,4,5),G(6,4,5),G(6,4,5),L(5,6,4),M(6,4),G(6),3,0.04,0.2,0.3,True,'k3'),
                         (G(6,4,5),G(6,4,5),G(6,4,5),L(5,6,4),M(6,4),G(6),3,0.04,0.2,0.3,False,'k1')])
    T("ppo_objective", [(G(6),G(6),G(6),G(6),G(6),torch.tensor(0.3),0.99,0.95,0.2,0.2,0.2,0.5)])
    T("rloo_objective", [(G(6,4,5),G(6,4,5),L(5,6,4),M(6,4),G(6),3,0.2,0.3)])
    T("reverse_kl_objective", [(G(3,4,5),G(3,4,5),L(5,3,4),M(3,4),G(3),0.1)])
    T("top_p_mask", [(torch.softmax(G(2,6),-1),0.7),(torch.softmax(G(3,5),-1),0.9)])
    T("top_k_mask", [(G(2,6),3),(G(3,5),2)])
    T("argmax_tokens", [(G(3,7),),(G(2,4,5),)])
    T("mode_label", [(L(4,10),),(torch.tensor([1,1,2,2]),),(torch.tensor([3,3,1,1,2]),)])
    T("median_reward", [(G(4),),(G(5),)])
    T("quantile_lower", [(G(10),0.3),(G(8),0.75)])
    T("pad_mask_from_lengths", [(torch.tensor([2,3,0]),4),(torch.tensor([1,4,2]),5)])
    T("first_nonzero_index", [(M(3,4),),(torch.tensor([0.,0.,1.,0.]),)])
    T("cumulative_max", [(G(5),),(G(3,2),)])
    T("bucketize_reward", [(G(6),torch.tensor([-0.5,0.0,0.5,1.0]))])

nfail = sum(1 for v in results.values() if not v)
print(f"\n{len(results)-nfail}/{len(results)} functions pass")
if nfail:
    print("FAILING:", [k for k,v in results.items() if not v])

Run the checker against reforacle

shell
python3 check.py
output
49/49 functions pass

[stdout]
49/49 functions pass

Re-run checker with different seeds

shell
for s in 1 2 3; do python3 -c "
import torch; torch.manual_seed($s)
import check
" 2>&1 | tail -3; echo "---seed $s done---"; done
output
49/49 functions pass
---seed 1 done---

49/49 functions pass
---seed 2 done---

49/49 functions pass
---seed 3 done---

[stdout]
49/49 functions pass
---seed 1 done---

49/49 functions pass
---seed 2 done---

49/49 functions pass
---seed 3 done---

Edge case and batched stress test

shell
python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(7)
def G(*s): return torch.randn(*s)
def M(*s): return (torch.rand(*s)>0.4).float()
def L(h,*s): return torch.randint(0,h,s)
def chk(n,a):
    try: r=getattr(R,n)(*a)
    except Exception as e: print(n,'REF-ERR',str(e)[:50]); return
    m=getattr(O,n)(*a)
    r=torch.as_tensor(r); m=torch.as_tensor(m)
    ok = (r.shape==m.shape) and torch.allclose(r.float(),m.float(),atol=1e-5,rtol=1e-4,equal_nan=True)
    print(n, 'OK' if ok else f'FAIL ref={r} mine={m}')
# batched gae/returns
chk('gae',(G(2,5),G(2,5),torch.tensor([0.5,0.2]),0.9,0.95))
chk('discounted_returns',(G(2,5),0.9))
chk('lambda_returns',(G(3,4),G(3,4),torch.tensor([0.1,0.2,0.3]),0.99,0.9))
# masked_mean 3D axis tuple
chk('masked_mean',(G(2,3,4),M(2,3,4),(1,2)))
chk('masked_sum',(G(2,3,4),M(2,3,4),-1))
# huber delta<1 large residual
chk('huber_value_loss',(G(6)*5,G(6)*5,0.5))
# cross_entropy all valid, label_smoothing 0
chk('smoothed_nll',(G(5,8),L(8,5),0.3))
# top_p extreme
chk('top_p_mask',(torch.softmax(G(4,10),-1),0.99))
chk('top_p_mask',(torch.softmax(G(4,10),-1),0.01))
chk('top_k_mask',(G(3,8),1))
# quantile various
chk('quantile_lower',(G(20),0.0)); chk('quantile_lower',(G(20),1.0)); chk('quantile_lower',(G(20),0.5))
chk('median_reward',(G(6),))
# mode with unique tie all distinct
chk('mode_label',(torch.tensor([5,4,3,2,1]),))
# first_nonzero all zero row & 1d
chk('first_nonzero_index',(torch.zeros(3,4),))
chk('cumulative_max',(G(4,3),))
chk('bucketize_reward',(G(10),torch.tensor([-1.,0.,1.])))
chk('advantage_mean_std',(G(2,3,4),M(2,3,4)))
chk('sequence_logprob',(G(2,4,6),L(6,2,4),torch.zeros(2,4),True))
"
output
gae REF-ERR The size of tensor a (5) must match the size of te
discounted_returns FAIL ref=tensor([[ 0.8845, 0.7155, -1.0004, 0.7289, 3.0171],
        [-0.9941, 1.8150, -0.6028, 1.6148, 1.9302]]) mine=tensor([[ 0.8938, -0.9838, -0.0731, 0.4274, 1.2799],
        [ 2.5947, 3.9875, 2.4139, 3.3520, 1.9302]])
lambda_returns REF-ERR The size of tensor a (4) must match the size of te
masked_mean OK
masked_sum OK
huber_value_loss OK
smoothed_nll OK
top_p_mask OK
top_p_mask OK
top_k_mask OK
quantile_lower OK
quantile_lower OK
quantile_lower OK
median_reward OK
mode_label OK
first_nonzero_index OK
cumulative_max OK
bucketize_reward OK
advantage_mean_std OK
sequence_logprob OK

[stdout]
gae REF-ERR The size of tensor a (5) must match the size of te
discounted_returns FAIL ref=tensor([[ 0.8845, 0.7155, -1.0004, 0.7289, 3.0171],
        [-0.9941, 1.8150, -0.6028, 1.6148, 1.9302]]) mine=tensor([[ 0.8938, -0.9838, -0.0731, 0.4274, 1.2799],
        [ 2.5947, 3.9875, 2.4139, 3.3520, 1.9302]])
lambda_returns REF-ERR The size of tensor a (4) must match the size of te
masked_mean OK
masked_sum OK
huber_value_loss OK
smoothed_nll OK
top_p_mask OK
top_p_mask OK
top_k_mask OK
quantile_lower OK
quantile_lower OK
quantile_lower OK
median_reward OK
mode_label OK
first_nonzero_index OK
cumulative_max OK
bucketize_reward OK
advantage_mean_std OK
sequence_logprob OK

Confirm dim0 time convention for gae/discounted/lambda

shell
python3 -c "
import torch, reforacle as R
torch.manual_seed(7)
def G(*s): return torch.randn(*s)
r=G(2,5)
ref=R.discounted_returns(r,0.9)
# dim0 accumulation
T=r.shape[0]; out=torch.zeros_like(r); run=torch.zeros_like(r[0])
for t in reversed(range(T)): run=r[t]+0.9*run; out[t]=run
print('dim0 match', torch.allclose(ref,out,atol=1e-6))
# gae dim0
rew=G(4,3); val=G(4,3); nv=G(3); gamma,lam=0.9,0.95
gref=R.gae(rew,val,nv,gamma,lam)
adv=torch.zeros_like(rew); last=torch.zeros_like(rew[0])
for t in reversed(range(4)):
    vn=val[t+1] if t+1<4 else nv
    d=rew[t]+gamma*vn-val[t]; last=d+gamma*lam*last; adv[t]=last
print('gae dim0 match', torch.allclose(gref,adv,atol=1e-6))
lref=R.lambda_returns(rew,val,nv,gamma,lam)
print('lambda dim0', torch.allclose(lref,adv+val,atol=1e-6))
# 1D still time dim0 trivially
"
output
dim0 match True
gae dim0 match True
lambda dim0 True

[stdout]
dim0 match True
gae dim0 match True
lambda dim0 True

/app/objectives.py

− removed
def gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[-1]
    adv = torch.zeros_like(rewards)
    last = torch.zeros_like(rewards[..., 0])
    for t in reversed(range(T)):
        v_next = values[..., t + 1] if t + 1 < T else next_value
        delta = rewards[..., t] + gamma * v_next - values[..., t]
        last = delta + gamma * lam * last
        adv[..., t] = last
    return adv
+ added
def gae(rewards, values, next_value, gamma, lam):
    T = rewards.shape[0]
    adv = torch.zeros_like(rewards)
    last = torch.zeros_like(rewards[0])
    for t in reversed(range(T)):
        v_next = values[t + 1] if t + 1 < T else next_value
        delta = rewards[t] + gamma * v_next - values[t]
        last = delta + gamma * lam * last
        adv[t] = last
    return adv

/app/objectives.py

− removed
def discounted_returns(rewards, gamma):
    T = rewards.shape[-1]
    out = torch.zeros_like(rewards)
    run = torch.zeros_like(rewards[..., 0])
    for t in reversed(range(T)):
        run = rewards[..., t] + gamma * run
        out[..., t] = run
    return out
+ added
def discounted_returns(rewards, gamma):
    T = rewards.shape[0]
    out = torch.zeros_like(rewards)
    run = torch.zeros_like(rewards[0])
    for t in reversed(range(T)):
        run = rewards[t] + gamma * run
        out[t] = run
    return out

Rerun full checker and 2D return checks

shell
python3 check.py 2>&1 | tail -2
python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(11)
def G(*s): return torch.randn(*s)
for a in [(G(2,5),G(2,5),G(5),0.9,0.95)]:
    print('gae2d', torch.allclose(R.gae(*a),O.gae(*a),atol=1e-6))
for a in [(G(3,4),0.9),(G(6),0.95)]:
    print('dret', torch.allclose(R.discounted_returns(*a),O.discounted_returns(*a),atol=1e-6))
for a in [(G(4,3),G(4,3),G(3),0.99,0.9)]:
    print('lam2d', torch.allclose(R.lambda_returns(*a),O.lambda_returns(*a),atol=1e-6))
"
output
49/49 functions pass
gae2d True
dret True
dret True
lam2d True

[stdout]
49/49 functions pass
gae2d True
dret True
dret True
lam2d True

Final randomized passes and boundary checks

shell
for s in 42 99; do python3 -c "import torch; torch.manual_seed($s); import check" 2>&1 | grep -E "pass|FAIL"; done
python3 -c "
import torch, reforacle as R, objectives as O
b=torch.tensor([0.0,0.5,1.0])
v=torch.tensor([0.0,0.5,1.0,-0.1,0.49,0.99,2.0])
print('bucketize exact-boundary', torch.equal(R.bucketize_reward(v,b),O.bucketize_reward(v,b)))
print('mode dtype', R.mode_label(torch.tensor([1,1,2])).dtype, O.mode_label(torch.tensor([1,1,2])).dtype)
print('first_nonzero dtype', R.first_nonzero_index(torch.tensor([0.,1.])).dtype, O.first_nonzero_index(torch.tensor([0.,1.])).dtype)
"
output
49/49 functions pass
49/49 functions pass
bucketize exact-boundary False
mode dtype torch.int64 torch.int64
first_nonzero dtype torch.int64 torch.int64

[stdout]
49/49 functions pass
49/49 functions pass
bucketize exact-boundary False
mode dtype torch.int64 torch.int64
first_nonzero dtype torch.int64 torch.int64

Determine bucketize right convention

shell
python3 -c "
import torch, reforacle as R
b=torch.tensor([0.0,0.5,1.0])
v=torch.tensor([-0.1,0.0,0.3,0.5,0.7,1.0,2.0])
print('ref  ', R.bucketize_reward(v,b).tolist())
print('right=False', torch.bucketize(v,b,right=False).tolist())
print('right=True ', torch.bucketize(v,b,right=True).tolist())
"
output
ref   [0, 1, 1, 2, 2, 3, 3]
right=False [0, 0, 1, 1, 2, 2, 3]
right=True  [0, 1, 1, 2, 2, 3, 3]

[stdout]
ref   [0, 1, 1, 2, 2, 3, 3]
right=False [0, 0, 1, 1, 2, 2, 3]
right=True  [0, 1, 1, 2, 2, 3, 3]

/app/objectives.py

− removed
    return torch.bucketize(rewards, boundaries)
+ added
    return torch.bucketize(rewards, boundaries, right=True)

Verify bucketize fix and full checker

shell
python3 -c "
import torch, reforacle as R, objectives as O
b=torch.tensor([0.0,0.5,1.0])
v=torch.tensor([0.0,0.5,1.0,-0.1,0.49,0.99,2.0])
print('bucketize exact', torch.equal(R.bucketize_reward(v,b),O.bucketize_reward(v,b)))
"
python3 check.py 2>&1 | tail -1
output
bucketize exact True
49/49 functions pass

[stdout]
bucketize exact True
49/49 functions pass

Verify standalone, no reforacle dependency

shell
grep -n "reforacle" /app/objectives.py || echo "NO reforacle dependency"
python3 -c "import ast; ast.parse(open('/app/objectives.py').read()); print('syntax OK')"
python3 -c "
import sys, types
# simulate sandbox without reforacle
sys.modules['reforacle']=None
import objectives
fns=[f for f in dir(objectives) if not f.startswith('_') and callable(getattr(objectives,f))]
print('importable, callables:', len([f for f in fns if f not in ('torch','F')]))
"
bash /app/timer.sh
output
NO reforacle dependency
syntax OK
importable, callables: 49
Remaining time (hours:minutes):
1:48

[stdout]
NO reforacle dependency
syntax OK
importable, callables: 49
Remaining time (hours:minutes):
1:48
All 49 functions are implemented and verified. Here's a summary of what I did. ## Approach Since `reforacle` is the only ground truth with no readable source, I recovered each function's exact numerical convention purely by **probing** , calling the reference with crafted inputs and matching candidate formulas , then built a checker (`check.py`) that compares my `objectives.py` against `reforacle` across many random seeds, shapes, and edge cases. ## Conventions recovered (the non-obvious ones) - **`masked_mean`**: `sum(v*m)/sum(m)`, returns **0** on an empty mask (not NaN). - **`whiten`** / **`masked_whiten`**: TRL convention , `(v-μ)/√(var+1e-8)` using *masked* biased stats; `shift_mean=True` gives zero-mean output, `False` adds μ back. `masked_whiten` additionally multiplies by the mask. - **`grpo_advantages`**: group-centered; when scaling, divides by **biased** std with `eps=1e-6`. - **`rloo_advantages`**: leave-one-out `r_i - (Σ - r_i)/(n-1)`. - **KL estimators**: `k1=d`, `k2=½d²`, `k3=e^{-d}-1+d`; `reverse_kl=k3`, `symmetric_kl=cosh(d)-1`. - **`top_p_mask`**: returns a *renormalized* nucleus distribution (not a boolean mask), keeping the crossing token. - **`cumulative_max`**: **lagged/exclusive** running max along **dim 0**. - **`gae`/`discounted_returns`/`lambda_returns`**: time axis is **dim 0**. - **`mode_label`**: ties broken toward the **largest** value. - **`median_reward`**: linear-interpolated quantile(0.5), not `torch.median`. - **`bucketize_reward`**: `right=True`. - **`masked_max`**: global reduction, returns 0 on empty mask. - Composites (`grpo_objective`, `ppo_objective`, `rloo_objective`, `reverse_kl_objective`, `dpo_sequence_loss`) reproduce the exact primitive composition (PPO whitens advantages but computes returns from raw GAE, etc.). The final `/app/objectives.py` depends only on `torch` (no `reforacle` import), passes all 49 functions across every randomized/edge test, and parses cleanly in a sandbox without the reference present.
Final diff , what the agent changed

No file changes, the agent left the workspace untouched.

Verifier output , the proof behind the reward
tests/test.sh
=== SyncVals verifier (post-training objectives, sealed oracle) ===

FUNCTIONS PASSED: 44/49  (partial_score=0.898)
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_cbd646ceeaba4a90. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_cbd646ceeaba4a90 · verifier authoritative; classifier explanatory.