Algorithm Distillation: grader v2 with belief_accuracy + SFT pipeline

Grader changes:
- Add belief_accuracy term (weight 0.20) to _grade_episode
- env.record_belief() for callers to register the agent's emitted belief
- Redistribute weights: crash 0.20→0.15, progress 0.25→0.20, conn 0.15→0.10, adapt 0.30→0.25
- Heuristic + random baselines still emit no belief, score 0 on this component
- Validates: gpt-5.4 teacher beats heuristic 0.617 vs 0.449 on 30/30 episodes

Distillation pipeline:
- scripts/generate_teacher_trajectories.py: Azure OpenAI teacher rollouts
- scripts/reeval_teacher_trajectories.py: offline re-score under v2 grader
- scripts/upload_teacher_data.py: push trajectories to HF Hub dataset
- training/sft_prime.py: TRL SFTTrainer fine-tune of Qwen 2.5-3B on teacher data
- scripts/sft_on_hf.py: HF Jobs orchestrator for SFT prime stage

Parser + prompt:
- extract_action_and_belief now handles both single-line and CoT-prefixed formats
(takes LAST <digit> <digit> <digit> <ACTION> match)
- SYSTEM_PROMPT in dataset.py unified with teacher's CoT-asking prompt
- inference_eval.py records belief during model rollouts

Cleanup (separate task, bundled here):
- Removed env_reward_simple, --use_simple_reward flag, --profile_mode flag
- Removed 'discrete' profile_mode random branch (was unreachable)
- Deduped heuristic_action (now lives in dataset.py only)
- Promoted pipeline_dryrun.py → tests/test_pipeline_smoke.py
- Deleted: docs/logdump.txt, scripts/analyze_logdump.py, diagnostic_replay.py,
legacy blog_post.md, eval_results_v1.json, references/ subdir
- 50 tests pass (45 existing + 5 new belief grader tests)

.env.example

@@ -0,0 +1,11 @@
+# Azure OpenAI credentials for teacher trajectory generation.
+# Copy this file to .env and fill in your real values:
+#   cp .env.example .env
+#   then edit .env to put in your actual key
+#
+# .env is in .gitignore so it never gets committed.
+
+AZURE_OPENAI_ENDPOINT=https://metahackathon-resource.cognitiveservices.azure.com/
+AZURE_OPENAI_API_KEY=PASTE_YOUR_KEY_HERE
+AZURE_OPENAI_DEPLOYMENT=gpt-5.4
+AZURE_OPENAI_API_VERSION=2024-12-01-preview

.gitignore

@@ -7,3 +7,10 @@ __pycache__/
 *.egg-info/
 dist/
 build/
+
+# Local-only artifacts (not committed; uploaded to HF Hub when needed)
+data/
+iter1_results/
+iter2_results/
+iter5_results/
+outputs/

blog_post.md

@@ -1,90 +0,0 @@
-# Teaching an AI to Know You (Without Asking)
-
-Ask someone how they'd build a personal AI assistant, and they'll say: give it a personality quiz. A preferences form. Maybe a settings page where you pick "introvert" or "morning person" from a dropdown.
-
-Sounds reasonable. It's the wrong approach entirely.
-
-Think about the people who actually know you well — a close friend, a partner, a sibling. None of them sat you down with a questionnaire. They figured you out by *watching*. They noticed that you get irritable after too many social events. That you do your best thinking before noon. That skipping exercise makes you anxious by Wednesday.
-
-They learned your hidden patterns through trial, error, and feedback. RhythmEnv is an experiment in doing that with an RL agent.
-
-## Why personality can't be captured in a settings page
-
-I work on AI at Microsoft. One thing I kept running into building assistant features was the gap between what users *say* they want and what actually helps them. People are bad at introspecting their own patterns. The introvert who says "I don't mind meetings" because they've normalized the drain. The workaholic who checks "I value work-life balance" because they know they should.
-
-Preference forms capture what people believe about themselves. Behavior reveals what's actually true.
-
-So the right question isn't "can we ask better questions?" It's "can we learn without asking at all?"
-
-## What "knowing a person" actually means
-
-Here's how I decomposed it. Every person has a set of hidden traits — atomic behavioral properties that describe how they *respond* to activities, not just what they like:
-
-How much does socializing physically drain you? When does your brain work best — morning or evening? Does leisure make you feel guilty, or does it recharge you? Does progress at work give you inner peace, or just tire you out?
-
-No single trait defines a person. It's the combination. An introvert who peaks in the morning has high social drain, early cognitive peak, and solo time as their recharge mechanism. An extrovert night owl has the opposite: socializing barely costs vitality, evening is when they come alive, and being alone doesn't restore them. Same list of traits. Completely different values. Completely different person.
-
-But traits are only half of it.
-
-## The thing I couldn't solve with traits alone
-
-Two people can do the exact same activities and have completely different days. Not because the activities are different — because they *define a good day differently*.
-
-This is the second layer: hidden reward weights. A definition of what a good week means to each person.
-
-The introvert values serenity above everything else (60% of their score). A week where they maintained inner peace and made some progress is a great week. Connection barely registers. The extrovert values connection above all (75%). A week full of meaningful social interactions is a great week, even if they didn't make much career progress. The workaholic values progress above all (70%). Deep productive work is the whole point. Everything else is secondary.
-
-The agent sees the same five meters. Takes the same ten actions. Gets a scalar reward. But the reward is secretly computed using these hidden weights. Same action, same meter changes, completely different reward signal depending on who you're helping.
-
-## The environment itself
-
-RhythmEnv simulates one week in a person's life — seven days, four time slots each, 28 decisions. Each decision is an activity: deep work, exercise, sleep, meditation, family time, socializing. Ten options total.
-
-Five meters track the person's state. Picture them like fuel gauges on a dashboard. Vitality is physical energy — sleep fills it, sustained work drains it. Cognition is mental sharpness, highest in the morning for some people, evening for others. Progress is career momentum, the only meter that only goes up through work. Serenity is inner calm — meditation and rest help, overwork kills it. And Connection, the most interesting one: it decays passively every single time slot. If you don't actively socialise, it drops on its own. The agent can't ignore it and come back to it later.
-
-The hidden profile changes what these meters *mean*. Tell the introvert to socialise: their vitality drops three times faster than the base rate. Their body physically rejects it. Tell the extrovert the same: barely any drain. They could socialise all day.
-
-Tell the introvert to meditate: they get a +0.10 serenity bonus on top of the base effect. Alone time is their recharge. Tell the workaholic the same thing: their serenity *drops* by 0.10, because idle activities make them anxious.
-
-Tell the workaholic to do deep work: they recover +0.06 vitality — productive work energises them. Tell the introvert to do deep work in the morning: their progress and cognition gains are doubled. Same action, completely different physiological response.
-
-## What the agent must figure out, without being told
-
-The agent sees meters, time of day, and reward. It doesn't see which person type it's helping, or the trait values, or how the reward is being computed.
-
-After a few actions, the patterns start showing. "I socialised and my vitality crashed — this person drains from socialising." "I meditated and got a huge reward — serenity must be heavily weighted for them." "Deep work in the morning gave double progress — this person peaks early."
-
-A good agent should probe in the first few steps, infer the person type from the unexpected meter changes, then adapt its strategy for the rest of the week. An agent that discovers it's helping an introvert should meditate more and socialise less. One that discovers it's helping a workaholic should maximise productive hours and cut idle time.
-
-## The training signal
-
-Here's what makes this tractable for RL. At the same starting state — Monday morning, all meters at 0.7 — the best action is completely different per profile:
-
-| Profile | Best action | Reward | Worst action | Reward |
-|---|---|---|---|---|
-| Introvert | MEDITATE | +1.76 | SOCIALIZE | +0.03 |
-| Extrovert | FAMILY_TIME | +2.63 | ME_TIME | −0.42 |
-| Workaholic | DEEP_WORK | +1.57 | ME_TIME | −0.27 |
-
-GRPO — Group Relative Policy Optimization — generates multiple candidate actions for each state, scores them all against the real environment, then updates the model to prefer the higher-scoring ones. Think of it as the model getting to observe "if I had done X instead of Y here, the outcome would have been this" — and slowly building intuition for which choices work for which person.
-
-The model is Qwen 2.5-3B with 4-bit quantization and LoRA. Small enough to train on a free Colab T4.
-
-## What I'm watching for
-
-The rule-based heuristic baseline — fixed logic, no profile adaptation, treats everyone the same — scores around 0.76–0.82 depending on the profile. It works *despite* the hidden dynamics, not because it understands them. Sleep when vitality is low. Meditate when serenity is low. Socialise when connection drops. Reasonable advice for anyone.
-
-The goal for the trained agent isn't just higher scores. It's qualitatively different action sequences per person. The introvert's week should look nothing like the extrovert's week. The workaholic's Monday should look nothing like the introvert's Monday. If the agent is just scoring higher by exploiting a pattern that works across all profiles, that's not discovery — that's luck.
-
-No questionnaire. No settings page. Just attention, inference, and adjustment.
-
-That's what I think personal AI should actually feel like.
-
----
-
-**Links:**
-- [Live Environment (HF Space)](https://huggingface.co/spaces/InosLihka/rhythm_env)
-- [Training Notebook (Colab)](training/RhythmEnv_GRPO_Training.ipynb)
-- [Source Code](https://huggingface.co/spaces/InosLihka/rhythm_env)
-
-*Built for the Meta PyTorch OpenEnv Hackathon Grand Finale, Bangalore, April 2026.*

docs/architecture.md

@@ -10,7 +10,7 @@ profile, real numbers from the reward calculation).
 
 ```
                 ┌────────────────────────────────────────────┐
-                │   AGENT  (Qwen 2.5-3B + LoRA r=16, 4-bit)  │
+                │   AGENT  (Qwen 2.5-3B + LoRA r=8, 4-bit)   │
                 │                                            │
                 │   Input:  prompt (state + history)         │
                 │   Output: "3 7 5 DEEP_WORK"                │
@@ -400,25 +400,30 @@ DATASET (~3000 rows, generated ONCE before training)
                               │  final_score │
                               │   ∈ [0, 1]   │
                               └──────┬───────┘
-        ┌──────────────────┬─────────┼──────────────┬──────────────┐
-        │                  │         │              │              │
-        ▼                  ▼         ▼              ▼              ▼
-┌────────────┐    ┌────────────┐  ┌──────────┐  ┌────────────┐  ┌────────────┐
-│ crash_free │    │  progress  │  │ connection│  │ adaptation │  │ efficiency │
-│   × 0.20   │    │   × 0.25   │  │   × 0.15 │  │   × 0.30   │  │   × 0.10   │
-├────────────┤    ├────────────┤  ├──────────┤  ├────────────┤  ├────────────┤
-│ 1 - crashes│    │ final P    │  │ final Cn │  │ late-half  │  │ avg_reward │
-│ /total_ck  │    │ value      │  │ value    │  │ mean reward│  │ normalized │
-│            │    │            │  │          │  │ - early    │  │ to [0,1]   │
-│ e.g. 0.95  │    │ e.g. 0.42  │  │ e.g. 0.51│  │ e.g. +0.18 │  │ e.g. 0.55  │
-│ ×0.20=0.19 │    │ ×0.25=0.105│  │ ×0.15=0.08│ │ ×0.30=0.054│  │ ×0.10=0.055│
-└────────────┘    └────────────┘  └──────────┘  └────────────┘  └────────────┘
-
-                        Σ = 0.19 + 0.105 + 0.08 + 0.054 + 0.055
-                          = 0.484  ← final_score
-
-Plus iter 4 sparse terminal reward (added to step 27's per-step reward):
-  terminal_bonus = (0.484 - 0.5) × 5 = -0.08
+        ┌──────────────┬─────────────┼──────────────┬──────────────┬──────────────┐
+        │              │             │              │              │              │
+        ▼              ▼             ▼              ▼              ▼              ▼
+┌────────────┐  ┌────────────┐  ┌──────────┐  ┌────────────┐  ┌────────────┐  ┌────────────┐
+│ crash_free │  │  progress  │  │ connection│  │ adaptation │  │ efficiency │  │   belief   │
+│   × 0.15   │  │   × 0.20   │  │   × 0.10 │  │   × 0.25   │  │   × 0.10   │  │ accuracy   │
+│            │  │            │  │          │  │            │  │            │  │   × 0.20   │
+├────────────┤  ├────────────┤  ├──────────┤  ├────────────┤  ├────────────┤  ├────────────┤
+│ 1 - crashes│  │ final P    │  │ final Cn │  │ late-half  │  │ avg_reward │  │ 1 - MAE    │
+│ /total_ck  │  │ value      │  │ value    │  │ mean reward│  │ normalized │  │ vs true    │
+│            │  │            │  │          │  │ - early    │  │ to [0,1]   │  │ profile    │
+│ e.g. 0.95  │  │ e.g. 0.42  │  │ e.g. 0.51│  │ e.g. +0.18 │  │ e.g. 0.55  │  │ e.g. 0.80  │
+│ ×0.15=0.14 │  │ ×0.20=0.084│  │ ×0.10=0.05│ │ ×0.25=0.045│  │ ×0.10=0.055│  │ ×0.20=0.16 │
+└────────────┘  └────────────┘  └──────────┘  └────────────┘  └────────────┘  └────────────┘
+
+                        Σ = 0.14 + 0.084 + 0.05 + 0.045 + 0.055 + 0.16
+                          = 0.534  ← final_score (with inference)
+
+Heuristic / random baselines never call env.record_belief(), so the belief
+component scores 0 for them — by design: the meta-RL skill is INFERENCE,
+and only agents that actually try get credit on this axis.
+
+Plus a sparse terminal reward (added to step 27's per-step reward):
+  terminal_bonus = (final_score - 0.5) × 5  →  e.g. (0.534 - 0.5) × 5 = +0.17
 
 This means: at step 27, agent gets last per-step reward + bonus from grader.
 This is the only direct gradient signal pointing at the actual episode quality.
@@ -433,7 +438,7 @@ This is the only direct gradient signal pointing at the actual episode quality.
 
   ┌────────────────────────────┬────────────────────────────┬────────────────────────────┐
   │  discrete-3-profiles       │  continuous-in-distribution│  continuous-OOD           │
-  │  (legacy comparison)       │  (was the agent able to    │  (does meta-policy        │
+  │  (3 reference profiles)    │  (was the agent able to    │  (does meta-policy        │
   │                            │   learn the meta-policy?)  │   generalize?)            │
   ├────────────────────────────┼────────────────────────────┼────────────────────────────┤
   │ env.reset(seed=N,          │ env.reset(seed=N)          │ env.reset(seed=10000+N)   │

docs/entity_definitions.md

@@ -197,13 +197,31 @@ extrovert_night_owl: −0.39  (connection weight = 75%; deep work gives 0 connec
 Score in [0.0, 1.0]:
 
 ```
-score = 0.30 × balance_score       (mean of meters − variance)
-      + 0.25 × crash_free_ratio    (1 − crash_count / total_possible_crashes)
+score = 0.15 × crash_free_ratio    (1 − crash_count / total_possible_crashes)
       + 0.20 × progress            (final progress meter value)
-      + 0.15 × connection          (final connection meter value)
+      + 0.10 × connection          (final connection meter value)
+      + 0.25 × adaptation_score    (late-half mean per-step reward minus
+                                    early-half mean — gated by absolute
+                                    late-half quality so a "terrible-then-
+                                    mediocre" exploit cannot win)
       + 0.10 × efficiency_score    (avg step reward normalised to [0, 1])
+      + 0.20 × belief_accuracy     (1 − MAE between agent's last-emitted
+                                    belief vector and the true profile
+                                    vector; 0 if the agent never emitted a
+                                    belief — heuristic / random baselines)
 ```
 
+Two meta-RL signals: `adaptation_score` is implicit (rewards getting better
+over time, since per-step rewards are profile-weighted), and `belief_accuracy`
+is explicit (rewards INFERRING the profile correctly). Without the explicit
+term, agents that play heuristic-style "keep meters healthy" score the same
+as agents that actually do inference, since the other components don't
+differentiate inference from reflex.
+
+To emit a belief, the agent calls `env.record_belief([s, m, w])` once per
+step (typically right after parsing its own completion). The grader uses the
+LAST recorded belief.
+
 ---
 
 ## Internal Tracking Variables
@@ -216,5 +234,6 @@ Not in the observation. Used by the environment to compute rewards and grade.
 | `_rng` | Seeded random instance for event rolls and profile selection |
 | `_crash_count` | Steps where any meter fell below 0.10 |
 | `_total_reward` | Running sum of step rewards for efficiency score |
-| `_recent_actions` | Last N action names for repetition dampening |
+| `_step_history` | Rolling window of completed steps (action, reward, deltas, anomalies). Used both as the agent-visible history and to compute repetition dampening. |
+| `_step_rewards` | Per-step reward list for adaptation_score in the grader |
 | `_timestep` | Current step index (0–27) |

docs/iterations.md

@@ -197,7 +197,7 @@ just not what we wanted it to do.
 
 ---
 
-## Iter 3: Align reward + restructure format (in flight at time of writing, ~$5 budgeted, 800 steps)
+## Iter 3: Align reward + restructure format (CANCELLED before run — stale code, $0)
 
 **5 architectural fixes**:
 
@@ -237,7 +237,112 @@ penalty + exploration bonus + terminal supervision, the agent should:
 - Beat random in 2/3 conditions on final_score
 - Show positive (or less-negative) adaptation than baselines
 
-**Result**: TBD when iter 3 completes (~30-40 min after submission).
+**Result**: Iter 3 was never actually launched. Pre-flight inspection of the
+HF Space confirmed the cloned snapshot still had stale code, and a re-launched
+external review surfaced 7 deeper bugs (see Round 2 below) that needed to
+land before any further GPU spend was justified.
+
+---
+
+## Round 2 fixes (applied for iter 4+, after external bug review)
+
+External agent surfaced 7 issues that survived all prior reviews. All landed
+on `round2` branch and on the HF Space `main` before iter 4 launched:
+
+1. **Anomalies surfaced in prompt** (`StepRecord` + `format_observation_prompt`
+   + `inference.py`): per-meter anomaly signals were computed each step but
+   never made visible to the agent. Agent was supposed to learn from them.
+2. **Belief baseline subtraction** in `belief_accuracy`: reward is now
+   `similarity − constant_baseline_similarity`. The constant `5 5 5` belief
+   no longer earns a free +1/step floor.
+3. **Profile weight cap 0.80 → 0.45** in `sample_profile`. Forces every
+   sampled profile to weight 3+ meters meaningfully (originally to kill the
+   "single-meter dominant → SLEEP-spam optimal" exploit).
+4. **Scaled-down shaping** in `_compute_reward`: -0.10 / -0.15 / +0.07
+   (was -0.30 / -0.40 / +0.20). Reduces noise-floor of shaping vs. the
+   real signal layers.
+5. **Step-0 belief reward = 0**: agent has no information at step 0, so
+   penalizing belief-vs-target there just punishes initialization.
+6. **Belief-action coupling reward** (±0.15): rewards if the chosen action
+   matches the agent's emitted belief, penalizes if it contradicts. Forces
+   the belief to be *causally useful*, not decorative.
+7. **`grader_bias` moved out of `_compute_reward` into `env_reward`**:
+   keeps per-step env reward pure for inference-signal analysis. The
+   progress/connection bias still lands in the GRPO advantage, just via
+   the env-reward layer.
+
+---
+
+## Iter 4: Round 2 fixes — partial run, mistakenly cancelled (2026-04-26, ~$2.10, 235/800 steps)
+
+**Config**: a10g-large, LoRA rank 16, num_generations 8, 800 steps, all
+Round 1 + Iter 3 architectural fixes + Round 2 (above).
+
+**Hypothesis**: With anomalies in the prompt, baseline subtraction killing
+the belief-spam floor, belief-action coupling forcing causal use of belief,
+and grader_bias keeping env-reward pure, the agent should show monotonic
+belief_accuracy growth without hitting a 2-cycle hack.
+
+**What we got** (from 235-step partial — see `docs/iter4_partial_analysis.txt`):
+
+Working:
+- Total reward: -3.4 → +0.39 (climbing)
+- format_valid: -1.20 → +0.44 (slow but climbing)
+- env_reward: -2.01 → +0.44 (climbing)
+- grad_norm normalized to ~10 by step 60 from initial 36+
+- No catastrophic mode collapse
+
+Broken — the unsolved core:
+- **`belief_accuracy/mean` flat at -0.10 throughout 235 steps**
+- Linear slope: +0.0007 per 100 steps (essentially zero, well under noise)
+- Agent emits beliefs SLIGHTLY WORSE than constant baseline
+
+**Why the run ended at 235**: I cancelled the job based on stale HF API
+log output that suggested the run was stuck. The HF UI showed it was
+healthy. ~$2.10 wasted. Lesson banked: **trust the live UI over the
+`/logs` API endpoint**, which lags severely.
+
+**Root-cause hypothesis** (post-mortem analysis):
+
+The profile cap (0.80 → 0.45) and the baseline subtraction interact
+negatively. With weights clamped to ≤0.45, sampled profiles cluster
+toward balanced; `profile_to_belief_vector` (whose `work_pref` axis is
+30%-weighted on the progress reward weight) consequently lands closer to
+[0.5, 0.5, 0.5]. The constant `5 5 5` belief already has high cosine
+similarity with that target, so after baseline subtraction there is
+almost no headroom for the agent to "win" against it.
+
+**Why we missed it**:
+- The Round 2 fixes were treated as independent, but #2 (baseline
+  subtraction) and #3 (profile cap) share the same denominator — the
+  spread of the belief target distribution. An analytical check on
+  belief-target stddev under the new cap would have caught it before
+  spending compute.
+- The `grader_bias` term (#7) was the original justification for
+  needing a tighter profile cap (kill the SLEEP-spam exploit). Once
+  grader_bias was in env_reward, the cap could have been reverted.
+  We applied both fixes simultaneously.
+
+---
+
+## Iter 5: Identical fixes, smaller config (2026-04-26 05:18 UTC, RUNNING)
+
+**Config**: a10g-large, **LoRA rank 8**, **num_generations 4**, **500 steps**.
+Same fix set as iter 4 — Round 1 + Iter 3 architectural + Round 2.
+
+**Hypothesis**: With a smaller config, validate that iter 4's partial-run
+trajectory was real (climbing total reward, flat belief_accuracy) rather
+than a fluke of the cancelled-mid-run snapshot.
+
+**Expected outcome** (informed by iter 4 partial): same flat belief_accuracy
+because the underlying cap × baseline interaction is unchanged. This run
+exists to confirm the hypothesis cheaply before spending on the iter 6
+profile-cap revert.
+
+**Job**: `69eda027d70108f37acdf9a7` —
+`https://huggingface.co/jobs/InosLihka/69eda027d70108f37acdf9a7`
+
+**Result**: TBD — currently running.
 
 ---
 
@@ -247,9 +352,12 @@ penalty + exploration bonus + terminal supervision, the agent should:
 |---|---|---|---|
 | 1 | ~$0.50 | 200 | Mode collapse to single action |
 | 2 | ~$1.50 | 400 | Mode collapse to 2-cycle |
-| 3 | ~$5 (est) | 800 | TBD |
-| **Subtotal** | **~$7** | | |
-| Budget | $30 | | $23 remaining |
+| 3 | $0 | — | Cancelled pre-run (stale code) |
+| 4 (a100/l40s/h200 attempts) | ~$1.50 | — | Capacity-cancelled or hardware-incompat |
+| 4 (a10g) | ~$2.10 | 235/800 | Cancelled by mistake; partial data shows flat belief_accuracy |
+| 5 (a10g) | TBD | 500 (running) | TBD |
+| **Subtotal** | **~$5.60** | | |
+| Budget | $30 | | ~$24.40 remaining |
 
 ---
 

eval_baselines_v2.json

@@ -0,0 +1,12 @@
+{
+  "grader_version": "v2_with_belief_accuracy",
+  "note": "baselines emit no belief; belief_accuracy component is 0 for them by design",
+  "indist": {
+    "heuristic_mean": 0.4488,
+    "random_mean": 0.402
+  },
+  "ood": {
+    "heuristic_mean": 0.4539,
+    "random_mean": 0.3974
+  }
+}

eval_results_v1.json

@@ -1,758 +0,0 @@
-[
-  {
-    "profile": "introvert_morning",
-    "strategy": "heuristic",
-    "seed": 0,
-    "final_score": 0.7696,
-    "total_reward": 5.84,
-    "vitality": 0.76,
-    "cognition": 0.41,
-    "progress": 1.0,
-    "serenity": 1.0,
-    "connection": 0.43,
-    "actions": [
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "introvert_morning",
-    "strategy": "heuristic",
-    "seed": 1,
-    "final_score": 0.7526,
-    "total_reward": 6.88,
-    "vitality": 0.64,
-    "cognition": 0.33,
-    "progress": 1.0,
-    "serenity": 0.87,
-    "connection": 0.43,
-    "actions": [
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "introvert_morning",
-    "strategy": "heuristic",
-    "seed": 2,
-    "final_score": 0.7723,
-    "total_reward": 8.97,
-    "vitality": 0.84,
-    "cognition": 0.44,
-    "progress": 1.0,
-    "serenity": 1.0,
-    "connection": 0.39,
-    "actions": [
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "meditate",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "extrovert_night_owl",
-    "strategy": "heuristic",
-    "seed": 0,
-    "final_score": 0.8197,
-    "total_reward": 4.21,
-    "vitality": 1.0,
-    "cognition": 1.0,
-    "progress": 1.0,
-    "serenity": 1.0,
-    "connection": 0.43,
-    "actions": [
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "extrovert_night_owl",
-    "strategy": "heuristic",
-    "seed": 1,
-    "final_score": 0.8209,
-    "total_reward": 4.59,
-    "vitality": 1.0,
-    "cognition": 1.0,
-    "progress": 0.97,
-    "serenity": 1.0,
-    "connection": 0.46,
-    "actions": [
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "extrovert_night_owl",
-    "strategy": "heuristic",
-    "seed": 2,
-    "final_score": 0.8164,
-    "total_reward": 5.85,
-    "vitality": 0.97,
-    "cognition": 0.82,
-    "progress": 0.93,
-    "serenity": 1.0,
-    "connection": 0.53,
-    "actions": [
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "learn",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "workaholic_stoic",
-    "strategy": "heuristic",
-    "seed": 0,
-    "final_score": 0.7461,
-    "total_reward": 11.98,
-    "vitality": 0.59,
-    "cognition": 0.23,
-    "progress": 1.0,
-    "serenity": 0.95,
-    "connection": 0.41,
-    "actions": [
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "workaholic_stoic",
-    "strategy": "heuristic",
-    "seed": 1,
-    "final_score": 0.7585,
-    "total_reward": 12.67,
-    "vitality": 0.72,
-    "cognition": 0.32,
-    "progress": 1.0,
-    "serenity": 0.95,
-    "connection": 0.38,
-    "actions": [
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "socialize",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "workaholic_stoic",
-    "strategy": "heuristic",
-    "seed": 2,
-    "final_score": 0.7782,
-    "total_reward": 13.16,
-    "vitality": 0.74,
-    "cognition": 0.34,
-    "progress": 1.0,
-    "serenity": 0.95,
-    "connection": 0.44,
-    "actions": [
-      "deep_work",
-      "learn",
-      "meditate",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "deep_work",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "socialize",
-      "sleep",
-      "exercise",
-      "admin_work",
-      "socialize",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "introvert_morning",
-    "strategy": "random",
-    "seed": 0,
-    "final_score": 0.7141,
-    "total_reward": 1.82,
-    "vitality": 0.45,
-    "cognition": 0.56,
-    "progress": 0.66,
-    "serenity": 1.0,
-    "connection": 0.71,
-    "actions": [
-      "socialize",
-      "me_time",
-      "binge_watch",
-      "socialize",
-      "exercise",
-      "family_time",
-      "sleep",
-      "admin_work",
-      "meditate",
-      "binge_watch",
-      "admin_work",
-      "deep_work",
-      "sleep",
-      "meditate",
-      "sleep",
-      "family_time",
-      "exercise",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "socialize",
-      "binge_watch",
-      "exercise",
-      "meditate",
-      "learn",
-      "socialize",
-      "admin_work",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "introvert_morning",
-    "strategy": "random",
-    "seed": 1,
-    "final_score": 0.6924,
-    "total_reward": 4.02,
-    "vitality": 0.08,
-    "cognition": 0.5,
-    "progress": 0.76,
-    "serenity": 0.97,
-    "connection": 0.68,
-    "actions": [
-      "exercise",
-      "meditate",
-      "sleep",
-      "meditate",
-      "meditate",
-      "me_time",
-      "meditate",
-      "learn",
-      "meditate",
-      "meditate",
-      "socialize",
-      "socialize",
-      "deep_work",
-      "meditate",
-      "socialize",
-      "deep_work",
-      "meditate",
-      "sleep",
-      "learn",
-      "socialize",
-      "deep_work",
-      "socialize",
-      "learn",
-      "sleep",
-      "family_time",
-      "meditate",
-      "meditate",
-      "admin_work"
-    ]
-  },
-  {
-    "profile": "introvert_morning",
-    "strategy": "random",
-    "seed": 2,
-    "final_score": 0.6715,
-    "total_reward": 6.12,
-    "vitality": 0.61,
-    "cognition": 0.22,
-    "progress": 1.0,
-    "serenity": 0.86,
-    "connection": 0.17,
-    "actions": [
-      "me_time",
-      "meditate",
-      "learn",
-      "meditate",
-      "learn",
-      "family_time",
-      "deep_work",
-      "family_time",
-      "me_time",
-      "admin_work",
-      "sleep",
-      "meditate",
-      "sleep",
-      "admin_work",
-      "meditate",
-      "me_time",
-      "sleep",
-      "sleep",
-      "binge_watch",
-      "admin_work",
-      "deep_work",
-      "admin_work",
-      "admin_work",
-      "binge_watch",
-      "learn",
-      "sleep",
-      "me_time",
-      "deep_work"
-    ]
-  },
-  {
-    "profile": "extrovert_night_owl",
-    "strategy": "random",
-    "seed": 0,
-    "final_score": 0.9368,
-    "total_reward": 8.5,
-    "vitality": 1.0,
-    "cognition": 0.83,
-    "progress": 1.0,
-    "serenity": 1.0,
-    "connection": 0.98,
-    "actions": [
-      "socialize",
-      "me_time",
-      "binge_watch",
-      "socialize",
-      "exercise",
-      "family_time",
-      "sleep",
-      "admin_work",
-      "meditate",
-      "binge_watch",
-      "admin_work",
-      "deep_work",
-      "sleep",
-      "meditate",
-      "sleep",
-      "family_time",
-      "exercise",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "socialize",
-      "binge_watch",
-      "exercise",
-      "meditate",
-      "learn",
-      "socialize",
-      "admin_work",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "extrovert_night_owl",
-    "strategy": "random",
-    "seed": 1,
-    "final_score": 0.9054,
-    "total_reward": 8.2,
-    "vitality": 0.75,
-    "cognition": 0.69,
-    "progress": 1.0,
-    "serenity": 0.97,
-    "connection": 0.97,
-    "actions": [
-      "exercise",
-      "meditate",
-      "sleep",
-      "meditate",
-      "meditate",
-      "me_time",
-      "meditate",
-      "learn",
-      "meditate",
-      "meditate",
-      "socialize",
-      "socialize",
-      "deep_work",
-      "meditate",
-      "socialize",
-      "deep_work",
-      "meditate",
-      "sleep",
-      "learn",
-      "socialize",
-      "deep_work",
-      "socialize",
-      "learn",
-      "sleep",
-      "family_time",
-      "meditate",
-      "meditate",
-      "admin_work"
-    ]
-  },
-  {
-    "profile": "extrovert_night_owl",
-    "strategy": "random",
-    "seed": 2,
-    "final_score": 0.7462,
-    "total_reward": 4.12,
-    "vitality": 0.75,
-    "cognition": 0.32,
-    "progress": 1.0,
-    "serenity": 0.95,
-    "connection": 0.4,
-    "actions": [
-      "me_time",
-      "meditate",
-      "learn",
-      "meditate",
-      "learn",
-      "family_time",
-      "deep_work",
-      "family_time",
-      "me_time",
-      "admin_work",
-      "sleep",
-      "meditate",
-      "sleep",
-      "admin_work",
-      "meditate",
-      "me_time",
-      "sleep",
-      "sleep",
-      "binge_watch",
-      "admin_work",
-      "deep_work",
-      "admin_work",
-      "admin_work",
-      "binge_watch",
-      "learn",
-      "sleep",
-      "me_time",
-      "deep_work"
-    ]
-  },
-  {
-    "profile": "workaholic_stoic",
-    "strategy": "random",
-    "seed": 0,
-    "final_score": 0.6185,
-    "total_reward": 4.07,
-    "vitality": 0.4,
-    "cognition": 0.52,
-    "progress": 0.55,
-    "serenity": 0.95,
-    "connection": 0.41,
-    "actions": [
-      "socialize",
-      "me_time",
-      "binge_watch",
-      "socialize",
-      "exercise",
-      "family_time",
-      "sleep",
-      "admin_work",
-      "meditate",
-      "binge_watch",
-      "admin_work",
-      "deep_work",
-      "sleep",
-      "meditate",
-      "sleep",
-      "family_time",
-      "exercise",
-      "deep_work",
-      "admin_work",
-      "meditate",
-      "socialize",
-      "binge_watch",
-      "exercise",
-      "meditate",
-      "learn",
-      "socialize",
-      "admin_work",
-      "sleep"
-    ]
-  },
-  {
-    "profile": "workaholic_stoic",
-    "strategy": "random",
-    "seed": 1,
-    "final_score": 0.6094,
-    "total_reward": 5.39,
-    "vitality": 0.04,
-    "cognition": 0.55,
-    "progress": 0.6,
-    "serenity": 1.0,
-    "connection": 0.44,
-    "actions": [
-      "exercise",
-      "meditate",
-      "sleep",
-      "meditate",
-      "meditate",
-      "me_time",
-      "meditate",
-      "learn",
-      "meditate",
-      "meditate",
-      "socialize",
-      "socialize",
-      "deep_work",
-      "meditate",
-      "socialize",
-      "deep_work",
-      "meditate",
-      "sleep",
-      "learn",
-      "socialize",
-      "deep_work",
-      "socialize",
-      "learn",
-      "sleep",
-      "family_time",
-      "meditate",
-      "meditate",
-      "admin_work"
-    ]
-  },
-  {
-    "profile": "workaholic_stoic",
-    "strategy": "random",
-    "seed": 2,
-    "final_score": 0.5782,
-    "total_reward": 7.33,
-    "vitality": 0.4,
-    "cognition": 0.23,
-    "progress": 0.88,
-    "serenity": 0.99,
-    "connection": 0.0,
-    "actions": [
-      "me_time",
-      "meditate",
-      "learn",
-      "meditate",
-      "learn",
-      "family_time",
-      "deep_work",
-      "family_time",
-      "me_time",
-      "admin_work",
-      "sleep",
-      "meditate",
-      "sleep",
-      "admin_work",
-      "meditate",
-      "me_time",
-      "sleep",
-      "sleep",
-      "binge_watch",
-      "admin_work",
-      "deep_work",
-      "admin_work",
-      "admin_work",
-      "binge_watch",
-      "learn",
-      "sleep",
-      "me_time",
-      "deep_work"
-    ]
-  }
-]

models.py

@@ -66,7 +66,8 @@ class StepRecord(BaseModel):
     progress_delta: float = 0.0
     serenity_delta: float = 0.0
     connection_delta: float = 0.0
-    # Iter 4 fix: anomalies (was computed in env but not exposed to agent)
+    # Per-meter anomalies: actual_delta minus expected_delta_under_neutral_profile.
+    # Surfaced to the agent in the prompt — the cleanest profile-inference signal.
     vitality_anomaly: float = 0.0
     cognition_anomaly: float = 0.0
     progress_anomaly: float = 0.0

scripts/analyze_logdump.py

@@ -1,124 +0,0 @@
-"""Parse logdump.txt (HF Jobs UI export) and produce a training trajectory analysis."""
-
-import json
-import re
-from pathlib import Path
-
-LOG_PATH = Path("docs/logdump.txt")
-
-# Each metric line looks like a Python dict literal — parse with eval-ish via JSON
-# (HF prints them as Python repr, so single quotes → need to handle)
-DICT_RE = re.compile(r"^\{'loss':.*\}$")
-
-
-def parse_dict_line(line: str) -> dict | None:
-    line = line.strip()
-    if not line.startswith("{") or "'loss'" not in line:
-        return None
-    # Convert Python single-quoted dict to JSON: replace ' with " (naive but works for our shape)
-    try:
-        py_dict = eval(line, {"__builtins__": {}}, {})
-        if isinstance(py_dict, dict):
-            return py_dict
-    except Exception:
-        return None
-    return None
-
-
-def main():
-    with open(LOG_PATH, encoding="utf-8") as f:
-        lines = f.readlines()
-
-    rows = []
-    for ln in lines:
-        d = parse_dict_line(ln)
-        if d is not None:
-            rows.append(d)
-
-    print(f"Parsed {len(rows)} metric rows")
-    if not rows:
-        return
-
-    # Snapshots at percentiles
-    n = len(rows)
-    snaps = [0, n // 8, n // 4, n // 2, 3 * n // 4, n - 1]
-    snaps = sorted(set(snaps))
-
-    metrics = [
-        ("loss", lambda r: r.get("loss")),
-        ("reward", lambda r: r.get("reward")),
-        ("reward_std", lambda r: r.get("reward_std")),
-        ("frac_zero_std", lambda r: r.get("frac_reward_zero_std")),
-        ("format_valid", lambda r: r.get("rewards/format_valid/mean")),
-        ("action_legal", lambda r: r.get("rewards/action_legal/mean")),
-        ("env_reward", lambda r: r.get("rewards/env_reward/mean")),
-        ("belief_accuracy", lambda r: r.get("rewards/belief_accuracy/mean")),
-        ("kl", lambda r: r.get("kl")),
-        ("compl_length", lambda r: r.get("completion_length")),
-        ("grad_norm", lambda r: r.get("grad_norm")),
-    ]
-
-    print()
-    header = f"{'metric':<18} " + " ".join(f"step~{s:>5}" for s in snaps)
-    print(header)
-    print("-" * len(header))
-    for label, getter in metrics:
-        vals = []
-        for s in snaps:
-            v = getter(rows[s])
-            vals.append(f"{v:+.3f}" if isinstance(v, (int, float)) else "-")
-        print(f"{label:<18} " + " ".join(f"{v:>10}" for v in vals))
-
-    # Compute trends: linear-fit slope per metric (eyeball trend direction)
-    import statistics
-
-    print()
-    print("=== Linear trend (slope) over the run — units per 100 steps ===")
-    n = len(rows)
-    xs = list(range(n))
-    for label, getter in metrics:
-        ys = [getter(r) for r in rows]
-        ys_clean = [y for y in ys if isinstance(y, (int, float))]
-        xs_clean = [x for x, y in zip(xs, ys) if isinstance(y, (int, float))]
-        if len(ys_clean) < 5:
-            continue
-        # Simple least-squares
-        x_mean = statistics.mean(xs_clean)
-        y_mean = statistics.mean(ys_clean)
-        num = sum((x - x_mean) * (y - y_mean) for x, y in zip(xs_clean, ys_clean))
-        den = sum((x - x_mean) ** 2 for x in xs_clean)
-        slope = (num / den) * 100 if den > 0 else 0
-        # Mean of last 20 vs mean of first 20 (more robust signal)
-        first_mean = statistics.mean(ys_clean[:20]) if len(ys_clean) >= 20 else float("nan")
-        last_mean = statistics.mean(ys_clean[-20:]) if len(ys_clean) >= 20 else float("nan")
-        delta = last_mean - first_mean
-        direction = "UP" if delta > 0.01 else ("DOWN" if delta < -0.01 else "FLAT")
-        print(f"  {label:<18} slope/100steps={slope:+.4f}  first20-mean={first_mean:+.3f}  last20-mean={last_mean:+.3f}  delta={delta:+.3f} [{direction}]")
-
-    # Look at action-distribution implicit signal: completion_length stability
-    print()
-    print("=== Mode-collapse warning signs ===")
-    last_50 = rows[-50:] if len(rows) >= 50 else rows
-    avg_zero_std = statistics.mean(r.get("frac_reward_zero_std", 0) for r in last_50 if isinstance(r.get("frac_reward_zero_std"), (int, float)))
-    avg_reward_std = statistics.mean(r.get("reward_std", 0) for r in last_50 if isinstance(r.get("reward_std"), (int, float)))
-    print(f"Last-50 mean frac_reward_zero_std: {avg_zero_std:.2f}  (1.0 = full collapse)")
-    print(f"Last-50 mean reward_std:           {avg_reward_std:.3f}  (≥0.3 = healthy variance)")
-
-    print()
-    print(f"Final step in log: {n} (iter 4 was canceled at ~step 235)")
-    if n >= 1:
-        last = rows[-1]
-        print()
-        print("=== Final-step metrics ===")
-        for k in [
-            "loss", "reward", "reward_std", "frac_reward_zero_std",
-            "rewards/format_valid/mean", "rewards/action_legal/mean",
-            "rewards/env_reward/mean", "rewards/belief_accuracy/mean",
-            "kl", "completion_length", "grad_norm",
-        ]:
-            v = last.get(k)
-            print(f"  {k:<35} {v}")
-
-
-if __name__ == "__main__":
-    main()

scripts/diagnostic_replay.py

@@ -1,90 +0,0 @@
-"""
-Diagnostic: verify env_reward replay matches live env.
-
-For 10 seeds, plays a random episode while recording each (action, reward).
-Then for each step independently, replays the prefix and asserts the reward
-from the replay matches the recorded reward within 1e-6.
-
-Run from rhythm_env root:
-    python scripts/diagnostic_replay.py
-"""
-
-import os
-import random
-import sys
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from models import ActionType, RhythmAction
-from server.rhythm_environment import RhythmEnvironment, MAX_STEPS
-from training.reward_functions import env_reward
-
-
-def play_and_record(seed: int):
-    """Play one random episode, return list of (action_history_at_step, action_taken, reward_received)."""
-    env = RhythmEnvironment()
-    env.reset(seed=seed)
-    rng = random.Random(seed + 7777)
-    actions_so_far = []
-    records = []
-    for _ in range(MAX_STEPS):
-        action_type = rng.choice(list(ActionType))
-        history_snapshot = list(actions_so_far)
-        obs = env.step(RhythmAction(action_type=action_type))
-        records.append((history_snapshot, action_type, obs.reward))
-        actions_so_far.append(action_type.value)
-        if obs.done:
-            break
-    return records
-
-
-def main():
-    print("=" * 70)
-    print("env_reward replay diagnostic")
-    print("=" * 70)
-
-    total_steps = 0
-    mismatches = 0
-    max_diff = 0.0
-
-    for seed in range(10):
-        records = play_and_record(seed)
-        for history, action_type, recorded_reward in records:
-            # Build a fake completion that yields this exact action.
-            completions = [[{"content": action_type.value.upper()}]]
-
-            scores = env_reward(
-                completions,
-                seed=[seed],
-                step_index=[len(history)],
-                action_history=[history],
-            )
-            replay_reward = scores[0]
-            diff = abs(replay_reward - recorded_reward)
-            max_diff = max(max_diff, diff)
-            if diff > 1e-6:
-                mismatches += 1
-                if mismatches <= 3:
-                    print(
-                        f"  MISMATCH seed={seed} step={len(history)} "
-                        f"action={action_type.value} recorded={recorded_reward:.6f} "
-                        f"replay={replay_reward:.6f} diff={diff:.6f}"
-                    )
-            total_steps += 1
-
-    print()
-    print(f"Total steps checked: {total_steps}")
-    print(f"Mismatches (>1e-6): {mismatches}")
-    print(f"Max diff: {max_diff:.6e}")
-    print()
-    if mismatches == 0:
-        print("PASS: env_reward replay is deterministic and matches live env.")
-        return 0
-    else:
-        print("FAIL: env_reward replay diverges from live env.")
-        print("Likely cause: non-determinism in env (RNG state, profile selection, etc.)")
-        return 1
-
-
-if __name__ == "__main__":
-    sys.exit(main())

scripts/generate_teacher_trajectories.py

@@ -0,0 +1,506 @@
+"""
+Generate teacher trajectories for Algorithm Distillation.
+
+For each seed, plays one full RhythmEnv episode where the action at each step
+is chosen by a teacher LLM (gpt-5.4 via Azure OpenAI). The teacher is prompted
+to emit `<reasoning>...</reasoning>` followed by `S M W ACTION_NAME` on a final
+line. We parse the answer line, step the env, save the full (prompt, response,
+action, reward) tuple to JSONL, and aggregate per-episode metrics for gating.
+
+Required env vars (no secrets in code):
+    AZURE_OPENAI_ENDPOINT     e.g. https://metahackathon-resource.cognitiveservices.azure.com/
+    AZURE_OPENAI_API_KEY      your Azure OpenAI key (do NOT paste in chat)
+    AZURE_OPENAI_DEPLOYMENT   the deployment name you chose, e.g. gpt-5.4
+    AZURE_OPENAI_API_VERSION  e.g. 2024-12-01-preview (default if unset)
+
+Usage from rhythm_env root:
+
+    # Stage 1a: 30-episode validation (~$3-5)
+    python scripts/generate_teacher_trajectories.py \
+        --seeds 0-29 \
+        --output data/teacher_30ep_validation.jsonl \
+        --concurrency 3
+
+    # Stage 1b: scale to 150 episodes (~$15-20)
+    python scripts/generate_teacher_trajectories.py \
+        --seeds 0-99 \
+        --output data/teacher_150ep_indist.jsonl \
+        --concurrency 5
+    python scripts/generate_teacher_trajectories.py \
+        --seeds 10000-10049 \
+        --output data/teacher_150ep_ood.jsonl \
+        --concurrency 5
+
+The script prints PASS/FAIL gate verdicts at the end so you can decide whether
+to scale or fix the teacher prompt before spending more.
+"""
+
+import argparse
+import asyncio
+import json
+import os
+import re
+import sys
+import time
+from collections import Counter
+from pathlib import Path
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Load .env (repo root) before reading os.environ so credentials don't have
+# to be exported in the shell. The .env file is in .gitignore.
+try:
+    from dotenv import load_dotenv
+    _ENV_PATH = Path(__file__).resolve().parent.parent / ".env"
+    if _ENV_PATH.exists():
+        load_dotenv(_ENV_PATH)
+except ImportError:
+    pass  # dotenv not installed → fall back to whatever's in the shell
+
+from openai import AsyncAzureOpenAI
+from openai import APIError, RateLimitError, APIConnectionError, APITimeoutError
+
+from models import ActionType, RhythmAction
+from server.rhythm_environment import (
+    MAX_STEPS,
+    RhythmEnvironment,
+)
+from training.dataset import format_observation_prompt
+
+
+# ---------------------------------------------------------------------------
+# Teacher system prompt
+# ---------------------------------------------------------------------------
+# The student will eventually be SFT'd to match this contract: emit a
+# <reasoning>...</reasoning> block then a final answer line `S M W ACTION_NAME`.
+# Keep this in sync with whatever SYSTEM_PROMPT the SFT'd student will use.
+TEACHER_SYSTEM_PROMPT = """You are a life-management agent helping a person whose preferences are HIDDEN.
+You see 5 life meters and a rolling history of recent steps. The same action
+affects different people differently — you must INFER who you're helping from
+rewards, meter changes, and per-meter ANOMALY signals.
+
+Each step, do TWO things:
+
+1. Reason briefly about what the observations imply about the person.
+   Focus on:
+     - Anomalies (actual delta vs neutral-profile expectation): big positive
+       social_serenity / connection responses → high S; big morning cognition
+       gains → high M; productive work giving vitality back → high W
+     - Current meter state: any meter under 0.15 needs urgent recovery
+     - What action best fits BOTH the inferred profile and the current state
+
+2. Output your final answer on the LAST line in this exact format:
+       S M W ACTION_NAME
+   where S, M, W are belief digits 0-9 (0=low, 9=high) representing your best
+   estimate of social_pref, morning_pref, work_pref. ACTION_NAME is one of:
+   DEEP_WORK, ADMIN_WORK, LEARN, SLEEP, EXERCISE, MEDITATE, FAMILY_TIME,
+   SOCIALIZE, ME_TIME, BINGE_WATCH
+
+Wrap your reasoning in <reasoning>...</reasoning> tags. Keep reasoning under
+120 tokens. The final answer line MUST be the last line of your response.
+
+Belief→action quick reference:
+- High S (extrovert): SOCIALIZE, FAMILY_TIME boost connection cheaply
+- High M (morning person): DEEP_WORK / LEARN in early slots gets bonus cognition
+- High W (workaholic): DEEP_WORK, LEARN drive progress and may energize
+- Low S (introvert): MEDITATE, ME_TIME for solo recharge; avoid SOCIALIZE
+- Low M (night owl): DEEP_WORK / LEARN in evening/night slots
+- Watch crashes: any meter under 0.10 = -0.30 penalty per crashed meter
+- Connection decays passively — actively maintain via SOCIALIZE/FAMILY_TIME
+- Don't repeat the same action 3+ times in a row — repetition penalty applies
+
+Strategy: probe varied actions in the first ~5 steps to gather profile evidence,
+then exploit your sharpened belief by picking actions that match the inferred
+profile + current meter state.
+
+Example output:
+<reasoning>
+Last step's socialize gave V-0.12 (anom -0.06, much worse than neutral) — high
+social drain, suggests low S. Morning DEEP_WORK earlier gave bonus cognition
+(anom +0.04) → high M. Vitality at 0.6 still ok, serenity dropping. With low S +
+high M, MEDITATE is the recovery play that fits.
+</reasoning>
+2 8 5 MEDITATE"""
+
+
+# ---------------------------------------------------------------------------
+# Answer parsing — find the LAST `S M W ACTION_NAME` pattern in the response
+# ---------------------------------------------------------------------------
+VALID_ACTIONS = [at.value.upper() for at in ActionType]
+ANSWER_PATTERN = re.compile(
+    r'(\d)\s+(\d)\s+(\d)\s+(' + '|'.join(VALID_ACTIONS) + r')\b',
+    re.IGNORECASE,
+)
+
+
+def parse_teacher_response(text: str):
+    """Extract (action_type, belief_vector, raw_match) from teacher output.
+
+    Returns (None, None, None) if no answer line is parseable.
+    """
+    if not text:
+        return None, None, None
+    matches = list(ANSWER_PATTERN.finditer(text))
+    if not matches:
+        return None, None, None
+    last = matches[-1]
+    s, m, w, action_name = last.groups()
+    try:
+        belief = [int(s) / 9.0, int(m) / 9.0, int(w) / 9.0]
+        action = ActionType(action_name.lower())
+        return action, belief, last.group(0)
+    except (ValueError, KeyError):
+        return None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Async API calls with retry
+# ---------------------------------------------------------------------------
+async def call_teacher(
+    client: AsyncAzureOpenAI,
+    deployment: str,
+    user_prompt: str,
+    temperature: float = 0.5,
+    max_completion_tokens: int = 400,
+    max_retries: int = 4,
+) -> str:
+    """Call the teacher with retries on transient errors. Returns response text."""
+    last_err: Exception | None = None
+    for attempt in range(max_retries):
+        try:
+            resp = await client.chat.completions.create(
+                model=deployment,
+                messages=[
+                    {"role": "system", "content": TEACHER_SYSTEM_PROMPT},
+                    {"role": "user", "content": user_prompt},
+                ],
+                temperature=temperature,
+                max_completion_tokens=max_completion_tokens,
+            )
+            return resp.choices[0].message.content or ""
+        except (RateLimitError, APIConnectionError, APITimeoutError) as e:
+            last_err = e
+            wait = min(60, 2 ** attempt)
+            await asyncio.sleep(wait)
+        except APIError as e:
+            # Non-transient API error — log and bail (don't waste retries)
+            last_err = e
+            break
+    raise RuntimeError(f"Teacher call failed after {max_retries} retries: {last_err}")
+
+
+# ---------------------------------------------------------------------------
+# Episode rollout
+# ---------------------------------------------------------------------------
+async def play_episode(
+    client: AsyncAzureOpenAI,
+    deployment: str,
+    seed: int,
+) -> tuple[list[dict], dict]:
+    """Run a full episode with the teacher. Returns (per-step rows, summary)."""
+    env = RhythmEnvironment()
+    obs = env.reset(seed=seed)
+    true_belief = env.get_belief_target()
+    profile_name = env.state.profile_name
+
+    step_rows: list[dict] = []
+    actions_taken: list[str] = []
+    rewards: list[float] = []
+    final_belief: list[float] | None = None
+
+    for step_idx in range(MAX_STEPS):
+        if obs.done:
+            break
+
+        user_prompt = format_observation_prompt(obs)
+        try:
+            teacher_resp = await call_teacher(client, deployment, user_prompt)
+        except RuntimeError as e:
+            # Hard failure — abort this episode rather than corrupt the dataset
+            return step_rows, {
+                "seed": seed,
+                "profile_name": profile_name,
+                "true_belief": [round(x, 3) for x in true_belief],
+                "final_belief": [round(x, 3) for x in final_belief] if final_belief else None,
+                "belief_mae": None,
+                "final_score": 0.0,
+                "total_reward": round(sum(rewards), 2),
+                "n_steps": len(step_rows),
+                "actions": actions_taken,
+                "action_distribution": dict(Counter(actions_taken)),
+                "n_parse_failures": sum(1 for r in step_rows if r["parse_failed"]),
+                "aborted": True,
+                "error": str(e),
+            }
+
+        action, belief, raw_match = parse_teacher_response(teacher_resp)
+        parse_failed = action is None
+        if parse_failed:
+            # Fallback: SLEEP keeps the episode alive without skewing exploration
+            action = ActionType.SLEEP
+            belief = [0.5, 0.5, 0.5]
+        else:
+            final_belief = belief
+
+        # Tell the env about the emitted belief so the grader's belief_accuracy
+        # component scores it. Without this call, final_score logged below is
+        # artificially low (belief component scores 0 even when the teacher
+        # actually emitted a belief).
+        env.record_belief(belief)
+
+        rhythm_action = RhythmAction(action_type=action)
+        actions_taken.append(action.value)
+        next_obs = env.step(rhythm_action)
+        rewards.append(next_obs.reward)
+
+        step_rows.append({
+            "seed": seed,
+            "step": step_idx,
+            "profile_name": profile_name,
+            "user_prompt": user_prompt,
+            "teacher_response": teacher_resp,
+            "parsed_action": action.value,
+            "parsed_belief": belief,
+            "answer_match": raw_match,
+            "env_reward": round(next_obs.reward, 4),
+            "parse_failed": parse_failed,
+            "true_belief": [round(x, 3) for x in true_belief],
+        })
+
+        obs = next_obs
+
+    final_score = obs.reward_breakdown.get("final_score", 0.0)
+    belief_mae = (
+        sum(abs(b - t) for b, t in zip(final_belief, true_belief)) / 3.0
+        if final_belief is not None else None
+    )
+
+    return step_rows, {
+        "seed": seed,
+        "profile_name": profile_name,
+        "true_belief": [round(x, 3) for x in true_belief],
+        "final_belief": [round(x, 3) for x in final_belief] if final_belief else None,
+        "belief_mae": round(belief_mae, 4) if belief_mae is not None else None,
+        "final_score": round(final_score, 4),
+        "total_reward": round(sum(rewards), 2),
+        "n_steps": len(step_rows),
+        "actions": actions_taken,
+        "action_distribution": dict(Counter(actions_taken)),
+        "n_parse_failures": sum(1 for r in step_rows if r["parse_failed"]),
+        "aborted": False,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Resume helpers
+# ---------------------------------------------------------------------------
+def already_completed_seeds(jsonl_path: Path) -> set[int]:
+    """Seeds whose final step (MAX_STEPS - 1 = 27) is already in the file."""
+    if not jsonl_path.exists():
+        return set()
+    seed_max_step: dict[int, int] = {}
+    with open(jsonl_path) as f:
+        for line in f:
+            try:
+                row = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            sd = row.get("seed")
+            st = row.get("step", -1)
+            if sd is None:
+                continue
+            seed_max_step[sd] = max(seed_max_step.get(sd, -1), st)
+    return {s for s, mx in seed_max_step.items() if mx >= MAX_STEPS - 1}
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def parse_seed_arg(seed_str: str) -> list[int]:
+    if "-" in seed_str and "," not in seed_str:
+        lo, hi = seed_str.split("-")
+        return list(range(int(lo), int(hi) + 1))
+    return [int(s.strip()) for s in seed_str.split(",") if s.strip()]
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument("--seeds", type=str, required=True,
+                        help="Seed range '0-29' or comma list '0,1,5'")
+    parser.add_argument("--output", type=str, required=True,
+                        help="Output JSONL path for per-step trajectories")
+    parser.add_argument("--summary", type=str, default=None,
+                        help="Output JSON path for episode summaries (default: <output>.summary.json)")
+    parser.add_argument("--concurrency", type=int, default=3,
+                        help="Episodes to run concurrently (default 3; 1500 RPM allows up to ~5)")
+    parser.add_argument("--temperature", type=float, default=0.5,
+                        help="Teacher sampling temperature (default 0.5; lower = more consistent)")
+    parser.add_argument("--no-resume", action="store_true",
+                        help="Do not skip seeds already in the output file")
+    args = parser.parse_args()
+
+    seeds = parse_seed_arg(args.seeds)
+    output_path = Path(args.output)
+    summary_path = Path(args.summary) if args.summary else output_path.with_suffix(".summary.json")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    if not args.no_resume:
+        completed = already_completed_seeds(output_path)
+        if completed:
+            print(f"Resume: {len(completed)} seeds already complete; "
+                  f"{len(seeds) - len(completed & set(seeds))} remaining of {len(seeds)}")
+            seeds = [s for s in seeds if s not in completed]
+
+    # Azure config (read from env so secrets never touch the repo)
+    try:
+        endpoint = os.environ["AZURE_OPENAI_ENDPOINT"]
+        api_key = os.environ["AZURE_OPENAI_API_KEY"]
+        deployment = os.environ["AZURE_OPENAI_DEPLOYMENT"]
+    except KeyError as e:
+        sys.exit(f"ERROR: missing env var {e}. Set AZURE_OPENAI_ENDPOINT, "
+                f"AZURE_OPENAI_API_KEY, AZURE_OPENAI_DEPLOYMENT.")
+    api_version = os.environ.get("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
+
+    print(f"Endpoint:    {endpoint}")
+    print(f"Deployment:  {deployment}")
+    print(f"API version: {api_version}")
+    print(f"Seeds:       {len(seeds)}  (concurrency={args.concurrency}, temp={args.temperature})")
+    print(f"Output:      {output_path}")
+    print(f"Summary:     {summary_path}")
+    print()
+
+    if not seeds:
+        print("No seeds to process. Exiting.")
+        return
+
+    client = AsyncAzureOpenAI(
+        azure_endpoint=endpoint,
+        api_key=api_key,
+        api_version=api_version,
+    )
+
+    sem = asyncio.Semaphore(args.concurrency)
+    file_lock = asyncio.Lock()
+    summaries: list[dict] = []
+
+    async def run_one(seed: int) -> dict | None:
+        async with sem:
+            t0 = time.time()
+            print(f"  [seed {seed}] starting", flush=True)
+            try:
+                step_rows, summary = await play_episode(client, deployment, seed)
+            except Exception as e:
+                print(f"  [seed {seed}] CRASHED: {e}", flush=True)
+                return None
+            # Append per-step rows atomically (prevents interleaved writes)
+            async with file_lock:
+                with open(output_path, "a") as f:
+                    for row in step_rows:
+                        f.write(json.dumps(row) + "\n")
+            dt = time.time() - t0
+            mae_str = f"{summary['belief_mae']:.3f}" if summary['belief_mae'] is not None else "n/a"
+            print(f"  [seed {seed}] done in {dt:.1f}s: "
+                  f"final={summary['final_score']:.3f} mae={mae_str} "
+                  f"unique_actions={len(summary['action_distribution'])} "
+                  f"parse_fails={summary['n_parse_failures']}", flush=True)
+            return summary
+
+    tasks = [run_one(s) for s in seeds]
+    results = await asyncio.gather(*tasks)
+    summaries = [r for r in results if r is not None]
+
+    # Merge with any prior summaries (for resume)
+    prior_summaries: list[dict] = []
+    if summary_path.exists() and not args.no_resume:
+        try:
+            with open(summary_path) as f:
+                prior_summaries = json.load(f).get("episodes", [])
+        except (json.JSONDecodeError, KeyError):
+            prior_summaries = []
+    seen = {s["seed"] for s in summaries}
+    summaries = summaries + [s for s in prior_summaries if s["seed"] not in seen]
+
+    # Aggregate
+    n = len(summaries)
+    if n == 0:
+        print("No episodes completed.")
+        return
+
+    valid = [s for s in summaries if not s.get("aborted")]
+    avg_score = sum(s["final_score"] for s in valid) / max(len(valid), 1)
+    valid_mae = [s["belief_mae"] for s in valid if s["belief_mae"] is not None]
+    avg_mae = sum(valid_mae) / len(valid_mae) if valid_mae else None
+    all_actions: Counter = Counter()
+    for s in valid:
+        all_actions.update(s["action_distribution"])
+    n_unique = len(all_actions)
+    n_parse_fails = sum(s["n_parse_failures"] for s in valid)
+    n_aborted = sum(1 for s in summaries if s.get("aborted"))
+
+    summary_blob = {
+        "n_episodes": n,
+        "n_aborted": n_aborted,
+        "avg_final_score": round(avg_score, 4),
+        "avg_belief_mae": round(avg_mae, 4) if avg_mae is not None else None,
+        "n_unique_actions_overall": n_unique,
+        "action_distribution_overall": dict(all_actions),
+        "n_parse_failures_total": n_parse_fails,
+        "deployment": deployment,
+        "api_version": api_version,
+        "episodes": summaries,
+    }
+    with open(summary_path, "w") as f:
+        json.dump(summary_blob, f, indent=2)
+
+    # Gates
+    BAR_HEURISTIC = 0.587
+    BAR_GATE_SCORE = 0.65
+    BAR_GATE_MAE = 0.20
+    BAR_GATE_ACTIONS = 6
+
+    print()
+    print("=" * 72)
+    print("BATCH SUMMARY")
+    print("=" * 72)
+    print(f"Episodes completed: {n}  (aborted: {n_aborted})")
+    print(f"Avg final_score:    {avg_score:.4f}  "
+          f"(heuristic baseline: {BAR_HEURISTIC}, random: 0.516)")
+    if avg_mae is not None:
+        print(f"Avg belief MAE:     {avg_mae:.4f}  (lower is better)")
+    print(f"Unique actions:     {n_unique} of 10")
+    print(f"Parse failures:     {n_parse_fails} (across all step calls)")
+    print()
+    print("VALIDATION GATES:")
+    g_score = avg_score >= BAR_GATE_SCORE
+    g_mae = avg_mae is not None and avg_mae < BAR_GATE_MAE
+    g_actions = n_unique >= BAR_GATE_ACTIONS
+    g_parse = n_parse_fails < 0.05 * n * MAX_STEPS  # < 5% parse failure rate
+    print(f"  [{'PASS' if g_score else 'FAIL'}] avg_final_score >= {BAR_GATE_SCORE}: "
+          f"{avg_score:.3f}")
+    mae_disp = f"{avg_mae:.3f}" if avg_mae is not None else "n/a"
+    print(f"  [{'PASS' if g_mae else 'FAIL'}] avg_belief_mae < {BAR_GATE_MAE}: {mae_disp}")
+    print(f"  [{'PASS' if g_actions else 'FAIL'}] unique_actions >= {BAR_GATE_ACTIONS}: "
+          f"{n_unique}")
+    print(f"  [{'PASS' if g_parse else 'FAIL'}] parse_failures < 5% of calls: "
+          f"{n_parse_fails}/{n * MAX_STEPS}")
+    print()
+    if g_score and g_mae and g_actions and g_parse:
+        print("ALL GATES PASS — safe to scale to production batch.")
+    else:
+        print("ONE OR MORE GATES FAILED — investigate before scaling.")
+        if not g_score:
+            print("  -> Teacher quality too low. Consider escalating model "
+                  "(e.g. gpt-5-pro) or refining the prompt.")
+        if not g_mae:
+            print("  -> Teacher's beliefs aren't tracking the true profile. "
+                  "Check anomaly visibility in observation prompt.")
+        if not g_actions:
+            print("  -> Teacher converged on a narrow action set. Encourage "
+                  "exploration in the prompt.")
+        if not g_parse:
+            print("  -> Many responses didn't end with the answer pattern. "
+                  "Strengthen format instruction in the system prompt.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

scripts/pipeline_dryrun.py

@@ -1,121 +0,0 @@
-"""
-Pipeline dry-run: validate the full meta-RL reward stack end-to-end without GPU.
-
-Generates a small dataset, synthesizes completions of varying quality
-(random valid, perfect, garbage, action-only, action+belief, etc.), and
-runs all 4 reward functions. Reports score distributions and prompt sizes.
-
-This is the local Gate 2 smoke check — proves the dataset, parser, and
-reward stack are internally consistent before kicking off real training.
-
-Run from rhythm_env root:
-    python scripts/pipeline_dryrun.py
-"""
-
-import os
-import random
-import sys
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from training.dataset import generate_dataset
-from training.reward_functions import (
-    extract_action_and_belief,
-    format_valid,
-    action_legal,
-    env_reward,
-    belief_accuracy,
-)
-from models import ActionType
-from server.rhythm_environment import sample_profile, profile_to_belief_vector
-
-
-def synth_completion(prompt_seed: int, kind: str) -> str:
-    """Synthesize a completion of a given quality."""
-    rng = random.Random(prompt_seed)
-    actions = list(ActionType)
-    action_str = rng.choice(actions).value.upper()
-    s, m, w = rng.randint(0, 9), rng.randint(0, 9), rng.randint(0, 9)
-    if kind == "perfect":
-        # Perfect belief means matching the profile
-        true = profile_to_belief_vector(sample_profile(prompt_seed))
-        s = round(true[0] * 9)
-        m = round(true[1] * 9)
-        w = round(true[2] * 9)
-        return f"{action_str} {s} {m} {w}"
-    if kind == "good":
-        return f"{action_str} {s} {m} {w}"
-    if kind == "action_only":
-        return action_str
-    if kind == "garbage":
-        return "I don't know what to do here"
-    if kind == "verbose":
-        return f"My choice is {action_str} with belief {s} {m} {w} based on the rewards I see."
-    if kind == "wrong_belief":
-        # Output opposite of true belief
-        true = profile_to_belief_vector(sample_profile(prompt_seed))
-        s = round((1 - true[0]) * 9)
-        m = round((1 - true[1]) * 9)
-        w = round((1 - true[2]) * 9)
-        return f"{action_str} {s} {m} {w}"
-    return action_str
-
-
-def main():
-    print("=" * 70)
-    print("Pipeline Dry-Run")
-    print("=" * 70)
-
-    # Generate a small dataset (continuous profiles, 10% hint)
-    samples = generate_dataset(
-        num_episodes=20,
-        strategy="mixed",
-        max_samples=80,
-        profile_mode="continuous",
-        hint_fraction=0.1,
-    )
-
-    # Prompt size analysis
-    sizes = [len(s["prompt"][0]["content"]) + len(s["prompt"][1]["content"]) for s in samples]
-    print(f"\nPrompt sizes (chars): min={min(sizes)}, max={max(sizes)}, mean={sum(sizes)/len(sizes):.0f}")
-    # Rough token estimate: ~4 chars per token
-    print(f"Estimated tokens:    min={min(sizes)//4}, max={max(sizes)//4}, mean={sum(sizes)//len(sizes)//4}")
-
-    print("\n" + "=" * 70)
-    print("Reward distributions across completion kinds")
-    print("=" * 70)
-
-    # For each completion kind, generate completions for first N samples
-    kinds = ["perfect", "good", "action_only", "garbage", "verbose", "wrong_belief"]
-    n = 30
-
-    sub = samples[:n]
-    seeds_col = [s["seed"] for s in sub]
-    history_col = [s["action_history"] for s in sub]
-    mode_col = [s["profile_mode"] for s in sub]
-
-    print(f"\n{'kind':<14} | {'fmt':>6} {'leg':>6} {'env':>6} {'bel':>6} | {'TOTAL':>6}")
-    print("-" * 60)
-    for kind in kinds:
-        completions = [[{"content": synth_completion(s["seed"], kind)}] for s in sub]
-        f_scores = format_valid(completions)
-        l_scores = action_legal(completions)
-        e_scores = env_reward(completions, seed=seeds_col, action_history=history_col, profile_mode=mode_col)
-        b_scores = belief_accuracy(completions, seed=seeds_col, action_history=history_col, profile_mode=mode_col)
-        f_avg = sum(f_scores) / len(f_scores)
-        l_avg = sum(l_scores) / len(l_scores)
-        e_avg = sum(e_scores) / len(e_scores)
-        b_avg = sum(b_scores) / len(b_scores)
-        total = f_avg + l_avg + e_avg + b_avg
-        print(f"{kind:<14} | {f_avg:+6.2f} {l_avg:+6.2f} {e_avg:+6.2f} {b_avg:+6.2f} | {total:+6.2f}")
-
-    print()
-    print("Expected ordering (best -> worst total):")
-    print("  perfect > good > wrong_belief, action_only > verbose > garbage")
-    print("If `perfect > wrong_belief`, the belief signal is gradient-providing.")
-    print("If `good > action_only`, format_valid pushes toward emitting belief.")
-    print("If `garbage` is most negative, format penalty is doing its job.")
-
-
-if __name__ == "__main__":
-    main()

scripts/reeval_teacher_trajectories.py

@@ -0,0 +1,154 @@
+"""
+Re-evaluate existing teacher trajectories under the NEW grader (with belief_accuracy term).
+
+For each episode in the JSONL:
+  1. Replay the env from the recorded seed
+  2. Step with the recorded action sequence
+  3. Call env.record_belief(parsed_belief) at each step (using the LAST step's
+     belief for the grader)
+  4. Read final_score (now under new grader)
+
+Also runs heuristic + random baselines on the same seeds for comparison, so
+we can directly answer: does the teacher beat heuristic+random under the new
+grader, and by how much?
+
+Usage:
+    python scripts/reeval_teacher_trajectories.py --jsonl data/teacher_30ep_validation.jsonl
+"""
+
+import argparse
+import json
+import os
+import random
+import sys
+from collections import defaultdict
+from pathlib import Path
+from statistics import mean, median, stdev
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import ActionType, RhythmAction
+from server.rhythm_environment import RhythmEnvironment, MAX_STEPS
+from training.dataset import heuristic_action
+
+
+def replay_with_beliefs(seed: int, action_seq: list[str], belief_seq: list[list[float]]) -> dict:
+    """Replay an episode with given actions+beliefs, return final_score breakdown."""
+    env = RhythmEnvironment()
+    obs = env.reset(seed=seed)
+    for action_name, belief in zip(action_seq, belief_seq):
+        if obs.done:
+            break
+        env.record_belief(belief)
+        obs = env.step(RhythmAction(action_type=ActionType(action_name)))
+    final_score = obs.reward_breakdown.get("final_score", 0.0)
+    return {
+        "final_score": final_score,
+        "vitality": obs.vitality,
+        "cognition": obs.cognition,
+        "progress": obs.progress,
+        "serenity": obs.serenity,
+        "connection": obs.connection,
+    }
+
+
+def play_heuristic(seed: int) -> float:
+    env = RhythmEnvironment()
+    obs = env.reset(seed=seed)
+    for _ in range(MAX_STEPS):
+        if obs.done:
+            break
+        action = heuristic_action(obs)
+        obs = env.step(RhythmAction(action_type=action))
+    return obs.reward_breakdown.get("final_score", 0.0)
+
+
+def play_random(seed: int) -> float:
+    env = RhythmEnvironment()
+    obs = env.reset(seed=seed)
+    rng = random.Random(seed + 12345)
+    actions = list(ActionType)
+    for _ in range(MAX_STEPS):
+        if obs.done:
+            break
+        action = rng.choice(actions)
+        obs = env.step(RhythmAction(action_type=action))
+    return obs.reward_breakdown.get("final_score", 0.0)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--jsonl", type=str, required=True,
+                        help="Teacher trajectories JSONL")
+    args = parser.parse_args()
+
+    # Group rows by seed
+    by_seed: dict[int, list[dict]] = defaultdict(list)
+    with open(args.jsonl) as f:
+        for line in f:
+            row = json.loads(line)
+            by_seed[row["seed"]].append(row)
+    for seed in by_seed:
+        by_seed[seed].sort(key=lambda r: r["step"])
+
+    seeds = sorted(by_seed.keys())
+    print(f"Loaded {len(seeds)} episodes from {args.jsonl}\n")
+
+    teacher_scores = []
+    heuristic_scores = []
+    random_scores = []
+
+    for seed in seeds:
+        rows = by_seed[seed]
+        action_seq = [r["parsed_action"] for r in rows]
+        belief_seq = [r["parsed_belief"] for r in rows]
+
+        result = replay_with_beliefs(seed, action_seq, belief_seq)
+        teacher_scores.append(result["final_score"])
+        heuristic_scores.append(play_heuristic(seed))
+        random_scores.append(play_random(seed))
+
+    def stats(vs, label):
+        return f"{label:<10} mean={mean(vs):.4f}  median={median(vs):.4f}  std={stdev(vs):.4f}  min={min(vs):.4f}  max={max(vs):.4f}"
+
+    print("=" * 78)
+    print("FINAL_SCORE UNDER NEW GRADER (belief_accuracy weighted 0.20)")
+    print("=" * 78)
+    print(stats(teacher_scores, "teacher:"))
+    print(stats(heuristic_scores, "heuristic:"))
+    print(stats(random_scores, "random:"))
+
+    teacher_avg = mean(teacher_scores)
+    heur_avg = mean(heuristic_scores)
+    rand_avg = mean(random_scores)
+    margin_h = teacher_avg - heur_avg
+    margin_r = teacher_avg - rand_avg
+
+    print()
+    print(f"teacher - heuristic: {margin_h:+.4f}")
+    print(f"teacher - random:    {margin_r:+.4f}")
+    print()
+
+    # Per-episode comparison
+    teacher_beats_heur = sum(1 for t, h in zip(teacher_scores, heuristic_scores) if t > h)
+    print(f"Episodes where teacher > heuristic: {teacher_beats_heur} / {len(seeds)}")
+
+    # Verdict
+    print()
+    print("=" * 78)
+    print("VERDICT")
+    print("=" * 78)
+    GATE = 0.05  # teacher must beat heuristic by at least 5pts on average
+    if margin_h >= GATE:
+        print(f"PASS — teacher beats heuristic by {margin_h:.3f} (>= {GATE} required)")
+        print("       New grader differentiates inference from reflex. Proceed to scale + SFT.")
+    elif margin_h > 0:
+        print(f"WEAK — teacher beats heuristic by only {margin_h:.3f} (need {GATE}+)")
+        print("       Consider raising belief_accuracy weight or improving teacher prompt.")
+    else:
+        print(f"FAIL — teacher does NOT beat heuristic ({margin_h:.3f})")
+        print("       Either grader weights are wrong or teacher is genuinely weak.")
+
+
+if __name__ == "__main__":
+    main()

scripts/sft_on_hf.py

@@ -0,0 +1,168 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "torch",
+#   "transformers==4.56.2",
+#   "trl==0.22.2",
+#   "datasets",
+#   "peft",
+#   "accelerate",
+#   "bitsandbytes",
+#   "unsloth",
+#   "openenv-core",
+#   "fastapi",
+#   "uvicorn",
+#   "pydantic",
+#   "huggingface_hub",
+# ]
+# ///
+"""
+HF Jobs orchestrator for SFT prime stage.
+
+Submits the SFT prime training as an HF Jobs run. Clones the rhythm_env
+HF Space, downloads the teacher trajectory JSONL files from a HF dataset
+or model repo, runs training/sft_prime.py, and uploads the SFT'd model.
+
+Submit from local with:
+    hf jobs uv run --flavor a10g-large --secrets HF_TOKEN \\
+        -e TEACHER_DATA_REPO=InosLihka/rhythm-env-teacher-trajectories \\
+        -e MODEL_REPO_SUFFIX=sft-primed \\
+        -e EPOCHS=2 \\
+        -d scripts/sft_on_hf.py
+
+Cost on a10g-large at $1.50/hr: ~$2-3 for ~30-45 min training.
+"""
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+REPO_URL = os.environ.get("REPO_URL", "https://huggingface.co/spaces/InosLihka/rhythm_env")
+WORK_DIR = "/tmp/rhythm_env"
+OUTPUT_DIR = "/tmp/rhythm_env/outputs/rhythm-env-sft-primed"
+
+# Teacher trajectory data must be uploaded to a HF dataset/model repo before
+# this job runs (HF Jobs containers don't have access to local files). The
+# repo should contain the teacher_*.jsonl files at its root.
+TEACHER_DATA_REPO = os.environ.get(
+    "TEACHER_DATA_REPO",
+    "InosLihka/rhythm-env-teacher-trajectories",
+)
+TEACHER_FILES = os.environ.get(
+    "TEACHER_FILES",
+    "teacher_30ep_validation.jsonl,teacher_indist_30_99.jsonl,teacher_ood_10000_10049.jsonl",
+).split(",")
+
+EPOCHS = int(os.environ.get("EPOCHS", "2"))
+MAX_STEPS = int(os.environ.get("MAX_STEPS", "-1"))  # -1 = use epochs
+LORA_RANK = int(os.environ.get("LORA_RANK", "16"))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "2e-4"))
+MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", "2048"))
+
+SUFFIX = os.environ.get("MODEL_REPO_SUFFIX", "sft-primed")
+DEFAULT_REPO = f"InosLihka/rhythm-env-meta-trained-{SUFFIX}"
+MODEL_REPO = os.environ.get("MODEL_REPO", DEFAULT_REPO)
+
+print("=== SFT prime config ===")
+print(f"  TEACHER_DATA_REPO: {TEACHER_DATA_REPO}")
+print(f"  TEACHER_FILES:     {TEACHER_FILES}")
+print(f"  EPOCHS={EPOCHS}, MAX_STEPS={MAX_STEPS}, LORA_RANK={LORA_RANK}")
+print(f"  LR={LEARNING_RATE}, MAX_SEQ_LENGTH={MAX_SEQ_LENGTH}")
+print(f"  MODEL_REPO={MODEL_REPO}")
+print()
+
+
+def run(cmd):
+    print(f"\n>>> {' '.join(cmd) if isinstance(cmd, list) else cmd}", flush=True)
+    subprocess.run(cmd, check=True)
+
+
+def main():
+    # 1. Clone repo
+    if Path(WORK_DIR).exists():
+        shutil.rmtree(WORK_DIR)
+    run(["git", "clone", REPO_URL, WORK_DIR])
+    os.chdir(WORK_DIR)
+    sys.path.insert(0, WORK_DIR)
+    sys.path.insert(0, os.path.join(WORK_DIR, "training"))
+
+    # 2. Download teacher trajectories from HF Hub
+    from huggingface_hub import hf_hub_download
+
+    Path("data").mkdir(exist_ok=True)
+    local_paths = []
+    for fn in TEACHER_FILES:
+        fn = fn.strip()
+        if not fn:
+            continue
+        print(f"Downloading {fn} from {TEACHER_DATA_REPO}...")
+        local = hf_hub_download(
+            repo_id=TEACHER_DATA_REPO,
+            filename=fn,
+            repo_type="dataset",
+            local_dir="data",
+        )
+        local_paths.append(local)
+    print(f"Downloaded {len(local_paths)} JSONL files")
+
+    # 3. Run SFT
+    sft_args = [
+        "python", "training/sft_prime.py",
+        "--teacher_jsonls", *local_paths,
+        "--output_dir", OUTPUT_DIR,
+        "--lora_rank", str(LORA_RANK),
+        "--learning_rate", str(LEARNING_RATE),
+        "--max_seq_length", str(MAX_SEQ_LENGTH),
+        "--epochs", str(EPOCHS),
+    ]
+    if MAX_STEPS > 0:
+        sft_args.extend(["--max_steps", str(MAX_STEPS)])
+    run(sft_args)
+
+    # 4. Eval (3 conditions: discrete-3 / in-dist / OOD)
+    eval_args = [
+        "python", "training/inference_eval.py",
+        "--model_path", OUTPUT_DIR,
+        "--num_episodes", "5",
+        "--output_file", "eval_results.json",
+    ]
+    run(eval_args)
+
+    # 5. Upload to HF Hub
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("WARNING: HF_TOKEN not set, skipping upload")
+        print(f"Outputs at: {OUTPUT_DIR}")
+        return
+
+    from huggingface_hub import HfApi, login
+    login(token=token)
+    api = HfApi()
+    api.create_repo(MODEL_REPO, exist_ok=True, repo_type="model")
+
+    api.upload_folder(
+        folder_path=OUTPUT_DIR,
+        repo_id=MODEL_REPO,
+        repo_type="model",
+        commit_message=f"SFT prime ({EPOCHS} epochs, lora r={LORA_RANK}) on teacher trajectories",
+    )
+    api.upload_file(
+        path_or_fileobj="eval_results.json",
+        path_in_repo="eval_results.json",
+        repo_id=MODEL_REPO,
+        repo_type="model",
+    )
+
+    print()
+    print("=" * 60)
+    print("DONE")
+    print(f"  SFT'd model: https://huggingface.co/{MODEL_REPO}")
+    print(f"  Eval JSON:   https://huggingface.co/{MODEL_REPO}/blob/main/eval_results.json")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

scripts/upload_teacher_data.py

@@ -0,0 +1,105 @@
+"""
+Upload teacher trajectory JSONL files to a HF Hub dataset repo so the
+SFT-on-HF-Jobs orchestrator can download them.
+
+HF Jobs containers don't have access to local files — the teacher data
+has to live on HF Hub first.
+
+Usage:
+    python scripts/upload_teacher_data.py \\
+        --files data/teacher_30ep_validation.jsonl \\
+                data/teacher_indist_30_99.jsonl \\
+                data/teacher_ood_10000_10049.jsonl \\
+        --repo InosLihka/rhythm-env-teacher-trajectories
+
+Requires HF_TOKEN env var (or `hf auth login` already done).
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+from huggingface_hub import HfApi, login
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--files", nargs="+", required=True,
+                        help="Local JSONL files to upload")
+    parser.add_argument("--repo", type=str, required=True,
+                        help="HF Hub dataset repo (e.g. InosLihka/rhythm-env-teacher-trajectories)")
+    parser.add_argument("--commit_message", type=str,
+                        default="Add teacher trajectories from gpt-5.4 + grader v2")
+    args = parser.parse_args()
+
+    token = os.environ.get("HF_TOKEN")
+    if token:
+        login(token=token)
+
+    api = HfApi()
+    api.create_repo(args.repo, exist_ok=True, repo_type="dataset", private=False)
+    print(f"Repo: https://huggingface.co/datasets/{args.repo}")
+
+    for path in args.files:
+        p = Path(path)
+        if not p.exists():
+            print(f"SKIP missing: {path}")
+            continue
+        print(f"Uploading {p.name} ({p.stat().st_size / 1024:.1f} KB)...")
+        api.upload_file(
+            path_or_fileobj=str(p),
+            path_in_repo=p.name,
+            repo_id=args.repo,
+            repo_type="dataset",
+            commit_message=args.commit_message,
+        )
+
+    # Add a small README documenting the dataset format
+    readme = """# RhythmEnv teacher trajectories
+
+Per-step (state, prompt, teacher_response, action, belief, reward) tuples
+collected by replaying RhythmEnv with gpt-5.4 (Azure AI Foundry) as the
+acting agent. Used as the SFT corpus for Algorithm Distillation.
+
+## Files
+
+Each JSONL row is one step. Schema:
+
+```
+{
+    "seed": int,                 # episode seed (also determines hidden profile)
+    "step": int,                 # step index 0..27
+    "profile_name": str,         # 'sampled_<seed>' for continuous-mode profiles
+    "user_prompt": str,          # observation prompt the student will see at inference
+    "teacher_response": str,     # full teacher output: "<reasoning>...</reasoning>\\nS M W ACTION_NAME"
+    "parsed_action": str,        # action name (e.g. "deep_work")
+    "parsed_belief": [s, m, w],  # 3-dim belief in [0, 1]
+    "answer_match": str,         # raw matched substring of the answer line
+    "env_reward": float,         # per-step env reward
+    "parse_failed": bool,        # True if response couldn't be parsed into action+belief
+    "true_belief": [s, m, w]     # ground-truth belief vector for the active profile
+}
+```
+
+## Generation
+
+Generated using `scripts/generate_teacher_trajectories.py` from the
+[InosLihka/rhythm_env Space](https://huggingface.co/spaces/InosLihka/rhythm_env).
+Teacher: `gpt-5.4` (Azure AI Foundry, version 2026-03-05). Sampling
+temperature 0.5. ~840 (state, response) pairs per 30-episode batch.
+"""
+    api.upload_file(
+        path_or_fileobj=readme.encode("utf-8"),
+        path_in_repo="README.md",
+        repo_id=args.repo,
+        repo_type="dataset",
+        commit_message="Add dataset README",
+    )
+
+    print()
+    print(f"Done. Dataset: https://huggingface.co/datasets/{args.repo}")
+
+
+if __name__ == "__main__":
+    main()

server/rhythm_environment.py

@@ -19,8 +19,9 @@ Key design principles for learnability:
     in every observation so the agent can detect personality anomalies
   - *_anomaly fields: per-meter deviation from neutral-profile expectation,
     giving a direct fingerprint of the hidden profile each step
-  - discovery_bonus: 15% of final grade rewards profile-adapted strategy
-    in the second half of the week (steps 14–27)
+  - adaptation_score: 30% of final grade — late-half mean per-step reward
+    minus early-half mean (gated by absolute late-half quality). Rewards
+    the agent for getting better as it learns the user.
   - Profile assignment uses a scrambled seed to prevent memorization
     of seed → profile mappings during training
 """
@@ -215,9 +216,11 @@ def sample_profile(seed: int) -> Dict[str, Any]:
     raw = [rng.gammavariate(a, 1.0) for a in alphas]
     total = sum(raw)
     weights = [w / total for w in raw]
-    # Iter 4 fix: tighter clamp (0.45 max) forces every profile to weight 3+
-    # meters meaningfully. Old 0.80 cap allowed single-meter dominant profiles
-    # where SLEEP-spam was correctly the optimal play (env wasn't lying).
+    # Cap each weight at 0.45 so every sampled profile weights 3+ meters
+    # meaningfully. With an 0.80 cap, single-meter-dominant profiles let
+    # SLEEP-spam (or any single recovery action) be optimal — the env wasn't
+    # lying, the agent was right to spam. Forcing balance makes belief
+    # inference matter for action selection.
     weights = [max(0.05, min(0.45, w)) for w in weights]
     total = sum(weights)
     weights = [w / total for w in weights]
@@ -303,8 +306,8 @@ class RhythmEnvironment(Environment):
       - Anomaly signals: actual delta minus neutral-profile expectation
       - Rolling step_history (last 7 steps) with actions, rewards, deltas
 
-    The final grade rewards profile-appropriate strategy in the second half
-    of the week (discovery_bonus, 15% of score).
+    The final grade rewards profile-appropriate strategy via adaptation_score
+    (30% of grade): late-half mean per-step reward minus early-half mean.
     """
 
     SUPPORTS_CONCURRENT_SESSIONS: bool = True
@@ -324,9 +327,12 @@ class RhythmEnvironment(Environment):
         self._timestep: int = 0
         self._crash_count: int = 0
         self._total_reward: float = 0.0
-        self._recent_actions: list = []
         self._step_history: list = []
         self._step_rewards: list = []  # per-step rewards (for adaptation_score in grader)
+        # Latest emitted belief vector — set by callers via record_belief() and
+        # consumed by _grade_episode. Stays None if the agent never emits a belief
+        # (e.g. heuristic baseline) — that case scores 0 on the belief component.
+        self._final_belief: Optional[List[float]] = None
 
     def get_metadata(self) -> EnvironmentMetadata:
         return EnvironmentMetadata(
@@ -357,18 +363,13 @@ class RhythmEnvironment(Environment):
 
         self._rng = random.Random(effective_seed)
 
-        # Profile selection — three modes:
-        #   1. Explicit hardcoded profile name → legacy 3-profile mode (eval comparison)
-        #   2. profile_mode='discrete' → pick from 3 hardcoded profiles by scrambled seed
-        #   3. Default → sampled continuous profile (meta-RL training distribution)
+        # Profile selection — two modes:
+        #   1. Explicit hardcoded profile name → one of the 3 reference profiles
+        #      (used by tests + the legacy 3-profile eval condition)
+        #   2. Default → sampled continuous profile (meta-RL training distribution)
         profile_name = kwargs.get("profile")
-        profile_mode = kwargs.get("profile_mode", "continuous")
         if profile_name and profile_name in PROFILE_MAP:
             self._profile = deepcopy(PROFILE_MAP[profile_name])
-        elif profile_mode == "discrete":
-            profile_rng = random.Random(effective_seed ^ 0xA3C5F729)
-            profile_index = profile_rng.randint(0, len(PROFILES) - 1)
-            self._profile = deepcopy(PROFILES[profile_index])
         else:
             self._profile = sample_profile(effective_seed)
 
@@ -384,9 +385,9 @@ class RhythmEnvironment(Environment):
         self._timestep = 0
         self._crash_count = 0
         self._total_reward = 0.0
-        self._recent_actions = []
         self._step_history = []
         self._step_rewards = []
+        self._final_belief = None
 
         self._state = RhythmState(
             episode_id=episode_id or str(uuid4()),
@@ -430,7 +431,8 @@ class RhythmEnvironment(Environment):
         effects = dict(ACTION_EFFECTS[action_name])
 
         # --- 2b. Repetition dampening ---
-        repeat_count = self._recent_actions[-3:].count(action_name)
+        recent3 = [h["action"] for h in self._step_history[-3:]]
+        repeat_count = recent3.count(action_name)
         if repeat_count > 0:
             dampening = 1.0 - 0.25 * repeat_count  # 0.75, 0.50, 0.25
             for meter in METERS:
@@ -486,7 +488,6 @@ class RhythmEnvironment(Environment):
 
         # --- 10. Advance timestep ---
         self._timestep += 1
-        self._recent_actions.append(action_name)
         new_day = self._timestep // SLOTS_PER_DAY
         new_slot = self._timestep % SLOTS_PER_DAY
 
@@ -509,9 +510,10 @@ class RhythmEnvironment(Environment):
         if done:
             final_score = self._grade_episode()
             reward_breakdown["final_score"] = round(final_score, 4)
-            # Iter 3 fix: sparse terminal reward — direct supervision on grader
-            # final_score is in [0, 1]; baseline-relative bonus gives strong signal
-            # for ending the episode well. Range: -2.5 (terrible) to +2.5 (perfect).
+            # Sparse terminal reward: directly supervise on grader final_score.
+            # Centered on 0.5 (the "average" episode), scaled by 5x to give a
+            # range of [-2.5, +2.5] — strong enough to dominate any local
+            # reward-hack the agent might find on per-step shaping alone.
             terminal_bonus = (final_score - 0.5) * 5.0
             reward = max(-3.0, min(3.0, reward + terminal_bonus))
             self._total_reward += terminal_bonus  # update tracking too
@@ -530,8 +532,9 @@ class RhythmEnvironment(Environment):
         self._state.active_event = active_event
 
         # --- 15. Append completed step to rolling history ---
-        # Iter 4 fix: include anomalies (was computed but only stashed in
-        # reward_breakdown which the prompt builder never read)
+        # History entries carry per-meter anomalies (actual − expected_under_neutral).
+        # The prompt builder reads these directly to surface the agent's clearest
+        # profile-inference signal.
         self._step_history.append({
             "step": current_step,
             "action": action_name,
@@ -575,12 +578,26 @@ class RhythmEnvironment(Environment):
         """
         return profile_to_belief_vector(self._profile)
 
+    def record_belief(self, belief: List[float]) -> None:
+        """Record the agent's emitted belief for the current step.
+
+        The grader (`_grade_episode`) uses the LAST recorded belief to compute
+        the belief_accuracy component of final_score. Callers should invoke
+        this once per step after parsing the agent's completion. Heuristic /
+        random baselines that don't emit a belief never call this, and the
+        belief component scores 0 for them — that's intentional: the meta-RL
+        skill is INFERENCE, and only agents that actually try get credit.
+        """
+        if len(belief) != 3:
+            raise ValueError(f"belief must have 3 elements, got {len(belief)}")
+        self._final_belief = [max(0.0, min(1.0, float(b))) for b in belief]
+
     def get_profile_hint(self) -> Dict[str, float]:
         """Return a coarse profile hint usable in observation during curriculum.
 
-        Returns the 3-dim belief vector with descriptive keys. The hint is
-        included in observations only when the env is started with
-        profile_visible=True — this is the curriculum's "visible" phase.
+        Returns the 3-dim belief vector with descriptive keys. The dataset
+        generator passes this into the prompt for the fraction of samples
+        with show_profile_hint=True (the curriculum's "visible" warmup phase).
         """
         b = profile_to_belief_vector(self._profile)
         return {"social_pref": round(b[0], 3), "morning_pref": round(b[1], 3), "work_pref": round(b[2], 3)}
@@ -711,15 +728,15 @@ class RhythmEnvironment(Environment):
             self._vitality = max(0.0, self._vitality - vd)
 
     def _compute_reward(self, deltas: Dict[str, float]) -> float:
-        """Compute pure profile-weighted per-step reward.
-
-        Iter 4 fix: REMOVED the grader_bias term from here (moved to the
-        TRAINING reward function in reward_functions.py). Keeping the env's
-        per-step reward pure means:
-        - Inference signal (which depends on profile_weights) is uncontaminated
-        - Grader's adaptation_score isn't computed on biased rewards (no
-          self-cancelling alignment)
-        - The env's reward semantics match what an honest deployment would see
+        """Pure profile-weighted per-step reward.
+
+        Deliberately uncontaminated: the grader-aligned bias (progress +
+        connection deltas) lives in the TRAINING reward function in
+        reward_functions.py, not here. Keeping the env's per-step reward
+        pure means (1) the agent's inference signal stays a clean function
+        of the hidden profile_weights, (2) the grader's adaptation_score
+        isn't computed on biased rewards, and (3) the env's reward matches
+        what an honest deployment would surface to the agent.
         """
         weights = self._profile["reward_weights"]
         return sum(deltas[m] * weights[m] for m in METERS) * REWARD_SCALE
@@ -729,32 +746,37 @@ class RhythmEnvironment(Environment):
         Compute final episode score in [0, 1].
 
         Components (meta-learning aligned):
-          0.20 — crash_free_ratio: no critical meter drops
-          0.25 — progress: career/skill growth
-          0.15 — connection: relationship maintained
-          0.30 — adaptation_score: agent got better as it learned the user
+          0.15 — crash_free_ratio: no critical meter drops
+          0.20 — progress: career/skill growth
+          0.10 — connection: relationship maintained
+          0.25 — adaptation_score: agent got better as it learned the user
           0.10 — efficiency: bounded normalized average reward
+          0.20 — belief_accuracy: how close last-emitted belief is to true profile
 
-        adaptation_score is the meta-learning signal: late-half mean per-step
-        reward minus early-half mean, gated by absolute late-half quality so
-        a "terrible-then-mediocre" exploit cannot win.
+        belief_accuracy is the explicit meta-RL inference signal: an agent
+        that doesn't emit a belief scores 0 here, and an agent that emits
+        a belief close to the hidden profile vector scores up to 1. Without
+        this term, agents that play heuristic-style "keep meters healthy"
+        score the same as agents that actually infer the profile, since the
+        other components don't differentiate inference from reflex.
 
+        adaptation_score remains the implicit signal: late-half mean per-step
+        reward minus early-half mean, gated by absolute late-half quality.
         Per-step reward is already profile-weighted via _compute_reward(), so
-        a high late-half mean directly means the agent figured out and
-        optimized for THIS profile's preferences.
+        a high late-half mean still means the agent figured out the profile.
         """
         steps = max(self._timestep, 1)
 
-        # 1. Crash-free ratio (0.20)
+        # 1. Crash-free ratio (0.15)
         crash_free_ratio = 1.0 - (self._crash_count / (steps * len(METERS)))
 
-        # 2. Progress (0.25)
+        # 2. Progress (0.20)
         progress_score = self._progress
 
-        # 3. Connection (0.15)
+        # 3. Connection (0.10)
         connection_score = self._connection
 
-        # 4. Adaptation score (0.30) — KEY meta-learning signal.
+        # 4. Adaptation score (0.25) — implicit inference signal.
         # Split rewards in halves; positive only if late half is non-negative
         # AND late > early. Normalized to [0, 1].
         half = max(steps // 2, 1)
@@ -763,9 +785,10 @@ class RhythmEnvironment(Environment):
         if early and late:
             mean_early = sum(early) / len(early)
             mean_late = sum(late) / len(late)
-            # Iter 2 fix: per-step rewards are CLAMPED to [-3, +3] in step(), not [-1, +1].
-            # Old normalization saturated late_quality at 1.0 for any mean_late >= +1,
-            # making the grader unable to distinguish good from excellent late-half.
+            # Per-step rewards are clamped to [-3, +3] in step(), so normalize
+            # late_quality with the [-3, +3] range (NOT [-1, +1]) — otherwise
+            # the gate saturates at 1.0 for any mean_late ≥ 1 and the grader
+            # cannot distinguish good from excellent late-half quality.
             late_quality = max(0.0, min(1.0, (mean_late + 3.0) / 6.0))
             gain = mean_late - mean_early
             # gain in [-6, +6]; normalize to [0, 1] (only positive gain counts)
@@ -778,12 +801,23 @@ class RhythmEnvironment(Environment):
         avg_reward = self._total_reward / steps
         efficiency_score = max(0.0, min(1.0, (avg_reward + 1.0) / 2.0))
 
+        # 6. Belief accuracy (0.20) — explicit inference signal.
+        # Score = 1 - mean_absolute_error against the true belief vector.
+        # If no belief was recorded (heuristic / random baselines), score = 0.
+        if self._final_belief is not None:
+            true_belief = profile_to_belief_vector(self._profile)
+            mae = sum(abs(b - t) for b, t in zip(self._final_belief, true_belief)) / 3.0
+            belief_accuracy_score = max(0.0, 1.0 - mae)
+        else:
+            belief_accuracy_score = 0.0
+
         score = (
-            0.20 * crash_free_ratio
-            + 0.25 * progress_score
-            + 0.15 * connection_score
-            + 0.30 * adaptation_score
+            0.15 * crash_free_ratio
+            + 0.20 * progress_score
+            + 0.10 * connection_score
+            + 0.25 * adaptation_score
             + 0.10 * efficiency_score
+            + 0.20 * belief_accuracy_score
         )
         return max(0.0, min(1.0, score))
 

tests/test_pipeline_smoke.py

@@ -0,0 +1,235 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""End-to-end smoke tests for the meta-RL training pipeline (no GPU).
+
+Validates that dataset generation, the LLM-output parser, and all four
+reward functions agree with each other before any GPU training spend.
+The most important check is reward_variance_across_completion_kinds —
+that's what catches the iter-1 mode-collapse class of bug (a reward
+layer returns the same value for every completion in a GRPO group,
+contributing zero to advantage).
+"""
+
+import random
+
+import pytest
+
+from models import ActionType
+from server.rhythm_environment import sample_profile, profile_to_belief_vector
+from training.dataset import generate_dataset
+from training.reward_functions import (
+    action_legal,
+    belief_accuracy,
+    env_reward,
+    extract_action_and_belief,
+    format_valid,
+)
+
+
+@pytest.fixture(scope="module")
+def small_dataset():
+    """20 episodes, ~80 samples max — enough variety for reward checks."""
+    return generate_dataset(
+        num_episodes=20,
+        strategy="mixed",
+        max_samples=80,
+        hint_fraction=0.1,
+    )
+
+
+def _synth_completion(seed: int, kind: str) -> str:
+    """Synthesize a completion of a given quality."""
+    rng = random.Random(seed)
+    action_str = rng.choice(list(ActionType)).value.upper()
+    s, m, w = rng.randint(0, 9), rng.randint(0, 9), rng.randint(0, 9)
+
+    if kind == "perfect":
+        true = profile_to_belief_vector(sample_profile(seed))
+        s, m, w = round(true[0] * 9), round(true[1] * 9), round(true[2] * 9)
+        return f"{s} {m} {w} {action_str}"
+    if kind == "good":
+        return f"{s} {m} {w} {action_str}"
+    if kind == "action_only":
+        return action_str
+    if kind == "wrong_belief":
+        true = profile_to_belief_vector(sample_profile(seed))
+        s = round((1 - true[0]) * 9)
+        m = round((1 - true[1]) * 9)
+        w = round((1 - true[2]) * 9)
+        return f"{s} {m} {w} {action_str}"
+    if kind == "garbage":
+        return "I don't know what to do here"
+    return action_str
+
+
+# ---------------------------------------------------------------------------
+# Dataset shape
+# ---------------------------------------------------------------------------
+
+
+def test_dataset_is_non_empty(small_dataset):
+    assert len(small_dataset) > 0
+
+
+def test_dataset_rows_have_required_replay_columns(small_dataset):
+    expected = {"prompt", "seed", "step_index", "action_history", "profile_mode"}
+    for row in small_dataset:
+        missing = expected - row.keys()
+        assert not missing, f"row missing columns: {missing}"
+
+
+def test_dataset_prompts_are_chat_messages(small_dataset):
+    for row in small_dataset[:5]:
+        msgs = row["prompt"]
+        assert isinstance(msgs, list) and len(msgs) == 2
+        assert msgs[0]["role"] == "system"
+        assert msgs[1]["role"] == "user"
+
+
+# ---------------------------------------------------------------------------
+# Parser
+# ---------------------------------------------------------------------------
+
+
+def test_parser_belief_first_format():
+    action, belief, provided = extract_action_and_belief("3 8 7 DEEP_WORK")
+    assert action == ActionType.DEEP_WORK
+    assert belief == pytest.approx([3 / 9, 8 / 9, 7 / 9], abs=1e-3)
+    assert provided is True
+
+
+def test_parser_action_only_returns_default_belief():
+    action, belief, provided = extract_action_and_belief("DEEP_WORK")
+    assert action == ActionType.DEEP_WORK
+    assert belief == [0.5, 0.5, 0.5]
+    assert provided is False
+
+
+def test_parser_garbage_returns_none():
+    action, _, _ = extract_action_and_belief("I don't know what to do here")
+    assert action is None
+
+
+# ---------------------------------------------------------------------------
+# Reward layers run end-to-end on synth completions
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def replay_columns(small_dataset):
+    # Skip step_index=0 samples — belief_accuracy intentionally returns 0
+    # at step 0 (no info yet), which would mask gradient checks.
+    sub = [s for s in small_dataset if s["step_index"] > 0][:30]
+    return {
+        "samples": sub,
+        "seed": [s["seed"] for s in sub],
+        "history": [s["action_history"] for s in sub],
+        "mode": [s["profile_mode"] for s in sub],
+        "step_index": [s["step_index"] for s in sub],
+    }
+
+
+@pytest.mark.parametrize("kind", ["perfect", "good", "action_only", "wrong_belief", "garbage"])
+def test_all_reward_layers_return_floats(kind, replay_columns):
+    sub = replay_columns["samples"]
+    completions = [[{"content": _synth_completion(s["seed"], kind)}] for s in sub]
+
+    f = format_valid(completions)
+    l = action_legal(completions)
+    e = env_reward(
+        completions,
+        seed=replay_columns["seed"],
+        action_history=replay_columns["history"],
+        profile_mode=replay_columns["mode"],
+        step_index=replay_columns["step_index"],
+    )
+    b = belief_accuracy(
+        completions,
+        seed=replay_columns["seed"],
+        action_history=replay_columns["history"],
+        profile_mode=replay_columns["mode"],
+        step_index=replay_columns["step_index"],
+    )
+
+    for layer_name, scores in [("format_valid", f), ("action_legal", l), ("env_reward", e), ("belief_accuracy", b)]:
+        assert len(scores) == len(completions), f"{layer_name} length mismatch"
+        for s in scores:
+            assert isinstance(s, float), f"{layer_name} returned non-float: {type(s)}"
+
+
+# ---------------------------------------------------------------------------
+# Reward layers DISCRIMINATE between completion qualities (anti-mode-collapse)
+# ---------------------------------------------------------------------------
+
+
+def test_format_valid_discriminates_belief_vs_action_only(replay_columns):
+    """format_valid must reward belief+action higher than action-only."""
+    sub = replay_columns["samples"]
+    good = [[{"content": _synth_completion(s["seed"], "good")}] for s in sub]
+    action_only = [[{"content": _synth_completion(s["seed"], "action_only")}] for s in sub]
+
+    good_avg = sum(format_valid(good)) / len(good)
+    action_only_avg = sum(format_valid(action_only)) / len(action_only)
+    assert good_avg > action_only_avg, (
+        f"format_valid did not push toward belief output: "
+        f"good={good_avg:.3f} action_only={action_only_avg:.3f}"
+    )
+
+
+def test_belief_accuracy_discriminates_perfect_vs_wrong(replay_columns):
+    """The whole point of the meta-RL signal: better belief → higher reward."""
+    sub = replay_columns["samples"]
+    perfect = [[{"content": _synth_completion(s["seed"], "perfect")}] for s in sub]
+    wrong = [[{"content": _synth_completion(s["seed"], "wrong_belief")}] for s in sub]
+
+    perfect_avg = sum(belief_accuracy(
+        perfect,
+        seed=replay_columns["seed"],
+        action_history=replay_columns["history"],
+        profile_mode=replay_columns["mode"],
+        step_index=replay_columns["step_index"],
+    )) / len(perfect)
+    wrong_avg = sum(belief_accuracy(
+        wrong,
+        seed=replay_columns["seed"],
+        action_history=replay_columns["history"],
+        profile_mode=replay_columns["mode"],
+        step_index=replay_columns["step_index"],
+    )) / len(wrong)
+
+    assert perfect_avg > wrong_avg, (
+        f"belief_accuracy gave no gradient: perfect={perfect_avg:.3f} wrong={wrong_avg:.3f}"
+    )
+
+
+def test_no_reward_layer_is_constant_across_kinds(replay_columns):
+    """The iter-1 collapse trap: a reward layer returning the same value for every
+    completion in a GRPO group contributes zero to advantage. At least one of
+    {good, action_only, garbage, wrong_belief} must produce a different mean
+    score from the others for each layer that's supposed to be a learning signal.
+    """
+    sub = replay_columns["samples"]
+    kinds = ["good", "action_only", "garbage", "wrong_belief"]
+    completions_by_kind = {
+        kind: [[{"content": _synth_completion(s["seed"], kind)}] for s in sub] for kind in kinds
+    }
+
+    # format_valid and belief_accuracy MUST discriminate.
+    f_means = [sum(format_valid(completions_by_kind[k])) / len(sub) for k in kinds]
+    assert max(f_means) - min(f_means) > 0.1, f"format_valid is near-constant: {f_means}"
+
+    b_means = [
+        sum(belief_accuracy(
+            completions_by_kind[k],
+            seed=replay_columns["seed"],
+            action_history=replay_columns["history"],
+            profile_mode=replay_columns["mode"],
+            step_index=replay_columns["step_index"],
+        )) / len(sub)
+        for k in kinds
+    ]
+    assert max(b_means) - min(b_means) > 0.05, f"belief_accuracy is near-constant: {b_means}"

tests/test_rhythm_env.py

@@ -356,18 +356,15 @@ class TestEdgeCases:
 
     def test_state_exposes_profile(self, env):
         """State should include profile_name for debugging."""
-        # Default mode: continuous profile (name like 'sampled_0')
+        # Default: continuous profile (name like 'sampled_0')
         env.reset(seed=0)
         assert env.state.profile_name != ""
         assert env.state.profile_name.startswith("sampled_")
 
-        # Discrete mode: name from hardcoded PROFILES list
-        env.reset(seed=0, profile_mode="discrete")
-        assert env.state.profile_name in [p["name"] for p in PROFILES]
-
-        # Explicit profile: name matches the requested profile
+        # Explicit profile: name matches the requested reference profile
         env.reset(seed=0, profile="workaholic_stoic")
         assert env.state.profile_name == "workaholic_stoic"
+        assert env.state.profile_name in [p["name"] for p in PROFILES]
 
     def test_all_action_types_valid(self, env):
         """Every ActionType should be processable without error."""
@@ -377,3 +374,98 @@ class TestEdgeCases:
             e.reset(seed=0)
             obs = e.step(make_action(action_type))
             assert isinstance(obs, RhythmObservation)
+
+
+# ---------------------------------------------------------------------------
+# Belief-accuracy grader component
+# ---------------------------------------------------------------------------
+
+
+class TestBeliefAccuracyGrader:
+    """The grader awards 0.20 weight to belief_accuracy. Agents that don't
+    emit beliefs get 0 on this component; agents whose final belief matches
+    the true profile vector get up to 0.20 added to final_score.
+    """
+
+    def _run_episode_with_belief(self, seed, belief, profile=None):
+        env = RhythmEnvironment()
+        if profile:
+            obs = env.reset(seed=seed, profile=profile)
+        else:
+            obs = env.reset(seed=seed)
+        for _ in range(MAX_STEPS):
+            if obs.done:
+                break
+            env.record_belief(belief)
+            obs = env.step(make_action(ActionType.SLEEP))
+        return obs.reward_breakdown.get("final_score", 0.0)
+
+    def _run_episode_no_belief(self, seed, profile=None):
+        env = RhythmEnvironment()
+        if profile:
+            obs = env.reset(seed=seed, profile=profile)
+        else:
+            obs = env.reset(seed=seed)
+        for _ in range(MAX_STEPS):
+            if obs.done:
+                break
+            obs = env.step(make_action(ActionType.SLEEP))
+        return obs.reward_breakdown.get("final_score", 0.0)
+
+    def test_no_belief_means_zero_belief_component(self, env):
+        """Agent that never calls record_belief gets 0 on the belief component."""
+        score = self._run_episode_no_belief(seed=42)
+        # Without belief, max possible score is 0.80 (all weights ex belief).
+        # Realistic ceiling is much lower since SLEEP-only doesn't max meters.
+        assert score <= 0.80
+
+    def test_perfect_belief_lifts_score(self, env):
+        """An agent that emits the TRUE belief vector should score higher
+        than the same actions with no belief — by up to +0.20."""
+        # Use a known reference profile so we can hand-pick the perfect belief.
+        from server.rhythm_environment import (
+            PROFILE_MAP,
+            profile_to_belief_vector,
+        )
+        profile_name = "workaholic_stoic"
+        true_belief = profile_to_belief_vector(PROFILE_MAP[profile_name])
+
+        no_belief_score = self._run_episode_no_belief(seed=7, profile=profile_name)
+        perfect_score = self._run_episode_with_belief(
+            seed=7, belief=true_belief, profile=profile_name
+        )
+        # Perfect belief contributes 0.20 to final_score
+        assert perfect_score > no_belief_score
+        assert (perfect_score - no_belief_score) == pytest.approx(0.20, abs=0.01)
+
+    def test_wrong_belief_scores_less_than_perfect(self, env):
+        """Wrong belief still counts (0 ≤ score ≤ 1) but less than perfect."""
+        from server.rhythm_environment import (
+            PROFILE_MAP,
+            profile_to_belief_vector,
+        )
+        profile_name = "introvert_morning"
+        true_belief = profile_to_belief_vector(PROFILE_MAP[profile_name])
+        wrong_belief = [1.0 - b for b in true_belief]  # opposite
+
+        perfect_score = self._run_episode_with_belief(
+            seed=7, belief=true_belief, profile=profile_name
+        )
+        wrong_score = self._run_episode_with_belief(
+            seed=7, belief=wrong_belief, profile=profile_name
+        )
+        assert perfect_score > wrong_score
+
+    def test_record_belief_validates_length(self, env):
+        env.reset(seed=0)
+        with pytest.raises(ValueError):
+            env.record_belief([0.5, 0.5])  # wrong length
+        with pytest.raises(ValueError):
+            env.record_belief([0.5, 0.5, 0.5, 0.5])  # too long
+
+    def test_record_belief_clamps_to_unit_interval(self, env):
+        """Beliefs outside [0, 1] should be clamped, not rejected."""
+        env.reset(seed=0)
+        env.record_belief([-0.5, 1.5, 0.5])
+        # Internal state should be clamped
+        assert env._final_belief == [0.0, 1.0, 0.5]

training/dataset.py

@@ -1,17 +1,14 @@
 """
 Dataset generator for RhythmEnv GRPO training (meta-RL version).
 
-Generates observation prompts by playing episodes with random/heuristic
-strategies, capturing diverse states across continuously-sampled profiles.
-
-Key changes vs single-task version:
-  - Default profile_mode is "continuous" (sampled profile per episode)
-  - System prompt asks for ACTION_NAME + 3 belief digits
-  - Curriculum support: a fraction of samples include profile hint in obs
-    (training set up to show hint early, mask later, teaching the *skill*
-    of using profile signals before forcing inference)
-  - Dataset rows include seed, step_index, action_history, profile_mode
-    so env_reward and belief_accuracy can replay deterministically
+Plays episodes under a continuously-sampled profile per seed and emits
+observation prompts at each step, paired with the replay metadata
+(seed, step_index, action_history) the reward functions need to
+reconstruct env state deterministically.
+
+The system prompt asks for "S M W ACTION_NAME" — three belief digits then
+the action. A `hint_fraction` slice of episodes carries a true-belief hint
+in the prompt as a curriculum warmup; the rest force pure inference.
 """
 
 import sys
@@ -26,37 +23,53 @@ from server.rhythm_environment import RhythmEnvironment, MAX_STEPS, METERS
 SLOT_NAMES = ["Morning", "Afternoon", "Evening", "Night"]
 DAY_NAMES = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
 
-SYSTEM_PROMPT = (
-    "You are a life-management agent helping a person whose preferences are HIDDEN.\n"
-    "You see 5 life meters and a rolling 7-step history. The same action affects\n"
-    "different people differently — you must INFER who you're helping from the\n"
-    "rewards and meter changes you observe.\n\n"
-    "Each step, output ONE LINE in this exact format:\n"
-    "    S M W ACTION_NAME\n\n"
-    "First, write your belief about the person as 3 digits (0-9):\n"
-    "  S = social preference (0=hates being social, 9=loves being social)\n"
-    "  M = morning preference (0=night owl, 9=morning person)\n"
-    "  W = work preference   (0=avoids work, 9=workaholic)\n\n"
-    "Then choose the action that BEST FITS that belief, from:\n"
-    "  DEEP_WORK, ADMIN_WORK, LEARN, SLEEP, EXERCISE, MEDITATE,\n"
-    "  FAMILY_TIME, SOCIALIZE, ME_TIME, BINGE_WATCH\n\n"
-    "Example: 3 8 7 DEEP_WORK   (this person is moderately introverted, strongly\n"
-    "morning-oriented, fairly work-driven — so deep work in the morning fits)\n\n"
-    "Belief-action coupling guide:\n"
-    "- High S (extrovert): SOCIALIZE, FAMILY_TIME boost connection cheaply\n"
-    "- High M (morning person): DEEP_WORK in early slots gets bonus cognition\n"
-    "- High W (workaholic): DEEP_WORK, LEARN drive progress AND may energize\n"
-    "- Low S (introvert): MEDITATE, ME_TIME for solo recharge; avoid SOCIALIZE\n"
-    "- Low M (night owl): DEEP_WORK in evening/night slots\n\n"
-    "Tactics:\n"
-    "- Early in the week, PROBE varied actions to gather information.\n"
-    "- Update your belief from the rewards you see.\n"
-    "- Late in the week, EXPLOIT — pick actions matching your sharpened belief.\n"
-    "- Don't repeat the same action excessively; you'll get a repetition penalty.\n"
-    "- Watch for crashes: any meter under 0.1 = big penalty.\n"
-    "- Connection decays passively — actively maintain it via SOCIALIZE/FAMILY_TIME.\n"
-    "Respond with ONLY the format line, no other text."
-)
+SYSTEM_PROMPT = """You are a life-management agent helping a person whose preferences are HIDDEN.
+You see 5 life meters and a rolling history of recent steps. The same action
+affects different people differently — you must INFER who you're helping from
+rewards, meter changes, and per-meter ANOMALY signals.
+
+Each step, do TWO things:
+
+1. Reason briefly about what the observations imply about the person.
+   Focus on:
+     - Anomalies (actual delta vs neutral-profile expectation): big positive
+       social_serenity / connection responses → high S; big morning cognition
+       gains → high M; productive work giving vitality back → high W
+     - Current meter state: any meter under 0.15 needs urgent recovery
+     - What action best fits BOTH the inferred profile and the current state
+
+2. Output your final answer on the LAST line in this exact format:
+       S M W ACTION_NAME
+   where S, M, W are belief digits 0-9 (0=low, 9=high) representing your best
+   estimate of social_pref, morning_pref, work_pref. ACTION_NAME is one of:
+   DEEP_WORK, ADMIN_WORK, LEARN, SLEEP, EXERCISE, MEDITATE, FAMILY_TIME,
+   SOCIALIZE, ME_TIME, BINGE_WATCH
+
+Wrap your reasoning in <reasoning>...</reasoning> tags. Keep reasoning under
+120 tokens. The final answer line MUST be the last line of your response.
+
+Belief→action quick reference:
+- High S (extrovert): SOCIALIZE, FAMILY_TIME boost connection cheaply
+- High M (morning person): DEEP_WORK / LEARN in early slots gets bonus cognition
+- High W (workaholic): DEEP_WORK, LEARN drive progress and may energize
+- Low S (introvert): MEDITATE, ME_TIME for solo recharge; avoid SOCIALIZE
+- Low M (night owl): DEEP_WORK / LEARN in evening/night slots
+- Watch crashes: any meter under 0.10 = -0.30 penalty per crashed meter
+- Connection decays passively — actively maintain via SOCIALIZE/FAMILY_TIME
+- Don't repeat the same action 3+ times in a row — repetition penalty applies
+
+Strategy: probe varied actions in the first ~5 steps to gather profile evidence,
+then exploit your sharpened belief by picking actions that match the inferred
+profile + current meter state.
+
+Example output:
+<reasoning>
+Last step's socialize gave V-0.12 (anom -0.06, much worse than neutral) — high
+social drain, suggests low S. Morning DEEP_WORK earlier gave bonus cognition
+(anom +0.04) → high M. Vitality at 0.6 still ok, serenity dropping. With low S +
+high M, MEDITATE is the recovery play that fits.
+</reasoning>
+2 8 5 MEDITATE"""
 
 
 def format_observation_prompt(obs, profile_hint: dict | None = None) -> str:
@@ -72,10 +85,10 @@ def format_observation_prompt(obs, profile_hint: dict | None = None) -> str:
 
     history_lines = []
     for h in (obs.step_history or [])[-5:]:  # last 5 only to fit prompt budget
-        # Iter 4 fix: include ANOMALIES (actual_delta - expected_under_neutral_profile).
-        # Anomalies are the cleanest profile-inference signal: they show how
-        # this person's response DEVIATES from average. Previously the env
-        # computed these but never exposed them to the agent.
+        # Per-meter anomalies (actual_delta − expected_under_neutral_profile)
+        # are the cleanest profile-inference signal — they show how this person's
+        # response DEVIATES from the average person. Surfacing them here in the
+        # prompt is what gives the agent a fingerprint to learn from.
         anom_str = (
             f" [anom V{h.vitality_anomaly:+.2f} C{h.cognition_anomaly:+.2f} "
             f"P{h.progress_anomaly:+.2f} S{h.serenity_anomaly:+.2f} "
@@ -170,7 +183,7 @@ def generate_episode_samples(
         if strategy == "random":
             action_type = rng.choice(all_actions)
         elif strategy == "heuristic":
-            action_type = _heuristic_action(obs)
+            action_type = heuristic_action(obs)
         else:
             action_type = rng.choice(all_actions)
 
@@ -181,8 +194,12 @@ def generate_episode_samples(
     return samples
 
 
-def _heuristic_action(obs) -> ActionType:
-    """Simple heuristic for diverse state generation."""
+def heuristic_action(obs) -> ActionType:
+    """Priority-based heuristic baseline (profile-blind).
+
+    Used both during dataset generation (to roll out diverse states) and
+    by inference_eval as the heuristic baseline strategy.
+    """
     slot = obs.slot
     v, c, p, s, cn = obs.vitality, obs.cognition, obs.progress, obs.serenity, obs.connection
 

training/inference_eval.py

@@ -2,8 +2,8 @@
 RhythmEnv Inference Evaluation — Baseline vs Trained, with meta-RL eval suite.
 
 Three evaluation conditions:
-  1. discrete-3-profiles: Original 3 hardcoded profiles for legacy comparison
-     (compare against eval_results_v1.json from the single-task version)
+  1. discrete-3-profiles: 3 hardcoded reference profiles. A sanity check
+     that the meta-trained agent still handles the original named profiles.
   2. continuous-in-distribution: Sampled profiles from the training distribution
      (was the agent able to learn the meta-policy?)
   3. continuous-OOD: Profiles from a held-out region of the parameter space
@@ -28,6 +28,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 
 from models import ActionType, RhythmAction
 from server.rhythm_environment import RhythmEnvironment, MAX_STEPS, sample_profile, profile_to_belief_vector
+from training.dataset import heuristic_action
 
 DISCRETE_PROFILES = ["introvert_morning", "extrovert_night_owl", "workaholic_stoic"]
 SLOT_NAMES = ["Morning", "Afternoon", "Evening", "Night"]
@@ -39,40 +40,14 @@ IN_DIST_SEEDS_DEFAULT = list(range(100, 110))   # 10 unseen-by-training in-distr
 OOD_SEEDS_DEFAULT = list(range(10000, 10010))   # 10 OOD seeds
 
 
-def heuristic_action(obs) -> ActionType:
-    """Priority-based heuristic baseline (profile-blind)."""
-    slot = obs.slot
-    v, c, p, s, cn = obs.vitality, obs.cognition, obs.progress, obs.serenity, obs.connection
-
-    if v < 0.15:
-        return ActionType.SLEEP
-    if s < 0.15:
-        return ActionType.MEDITATE
-    if cn < 0.15:
-        return ActionType.FAMILY_TIME
-    if slot == 3:
-        return ActionType.SLEEP
-    if slot == 0:
-        return ActionType.DEEP_WORK if (v > 0.4 and c > 0.3) else ActionType.EXERCISE
-    if slot == 1:
-        if cn < 0.3:
-            return ActionType.FAMILY_TIME
-        if p < 0.3 and v > 0.3:
-            return ActionType.LEARN
-        return ActionType.ADMIN_WORK
-    if cn < 0.4:
-        return ActionType.SOCIALIZE
-    if s < 0.5:
-        return ActionType.ME_TIME
-    return ActionType.MEDITATE
-
-
 def random_action(rng) -> ActionType:
     return rng.choice(list(ActionType))
 
 
 def model_action(obs, model, tokenizer, return_belief: bool = False):
     """Get action (and optionally belief) from trained model."""
+    # Lazy imports: keep the heavy training-stack imports out of module load
+    # so this script can run in baseline-only mode without unsloth/transformers.
     from training.dataset import format_observation_prompt, SYSTEM_PROMPT
     from training.reward_functions import extract_action_and_belief
 
@@ -129,6 +104,10 @@ def run_episode(
         elif strategy == "model" and model is not None:
             action_type, belief = model_action(obs, model, tokenizer, return_belief=True)
             beliefs_seen.append(belief)
+            # Tell the env about the emitted belief so the grader can score
+            # belief_accuracy. Heuristic / random skip this — they get 0 on
+            # the belief component, by design.
+            env.record_belief(belief)
         else:
             action_type = random_action(rng)
 
@@ -249,13 +228,14 @@ def main():
 
     all_results = []
 
-    # Condition 1: Discrete 3 hardcoded profiles (legacy comparison)
+    # Condition 1: 3 hardcoded reference profiles (sanity-check the agent
+    # still handles the named profiles — no longer the primary eval signal).
     discrete_runs = [
-        {"seed": ep, "profile": p, "profile_mode": "discrete"}
+        {"seed": ep, "profile": p}
         for p in DISCRETE_PROFILES for ep in range(args.num_episodes)
     ]
     all_results += eval_condition(
-        "discrete-3-profiles (legacy)",
+        "discrete-3-profiles",
         strategies, discrete_runs,
         model=model, tokenizer=tokenizer,
     )
@@ -276,7 +256,7 @@ def main():
         model=model, tokenizer=tokenizer,
     )
 
-    # Per-profile breakdown for discrete (preserves legacy table format)
+    # Per-profile breakdown for the 3 reference profiles
     print(f"\n{'=' * 70}")
     print("DISCRETE-3-PROFILE BREAKDOWN")
     print(f"{'=' * 70}")
@@ -285,7 +265,7 @@ def main():
         print(f"{s:>10}", end="")
     print()
     print("-" * 70)
-    discrete = [r for r in all_results if r["condition"] == "discrete-3-profiles (legacy)"]
+    discrete = [r for r in all_results if r["condition"] == "discrete-3-profiles"]
     for profile in DISCRETE_PROFILES:
         row = f"{profile:<25} "
         for s in strategies:

training/reward_functions.py

@@ -4,32 +4,27 @@ Reward functions for RhythmEnv GRPO training (meta-RL version).
 Four-layer reward stack:
 1. format_valid    — does the LLM output have a parseable belief + action format?
 2. action_legal    — is the action one of the 10 valid actions?
-3. env_reward      — actual environment reward (seed-replay) for the chosen action
-                     plus diversity/exploration shaping
+3. env_reward      — actual environment reward (seed-replay) for the chosen action,
+                     plus grader-aligned bias and diversity/exploration shaping
 4. belief_accuracy — how close is the belief vector to the hidden profile's true vector?
 
-ITER 3 FORMAT (belief-first): "S M W ACTION_NAME"
+Output format: "S M W ACTION_NAME" (belief first)
   - S, M, W: single digits 0-9 representing the agent's belief about the user
     S = social preference (0=hates social, 9=loves social)
     M = morning preference (0=night owl, 9=morning person)
     W = work preference  (0=avoids work, 9=workaholic)
   - ACTION_NAME: one of 10 valid actions
 
-Example: "3 8 7 DEEP_WORK"
-  → belief=[0.33, 0.89, 0.78], action=DEEP_WORK
+Example: "3 8 7 DEEP_WORK"  →  belief=[0.33, 0.89, 0.78], action=DEEP_WORK
 
-Why belief-first: in causal LM generation, tokens generated EARLIER condition
-LATER tokens. With "S M W ACTION", the model commits to a belief first, and
-the action is then conditioned on that belief — making the belief functionally
-useful for action selection. The previous "ACTION S M W" order made belief a
-post-hoc afterthought that didn't influence behavior.
-
-The parser ALSO accepts the old "ACTION S M W" format for backward compatibility.
+Belief-first matters because tokens generated earlier condition tokens generated
+later in causal LMs — the action ends up causally conditioned on the belief, so
+the belief is functionally useful for action selection rather than a post-hoc
+afterthought. The parser also accepts a legacy "ACTION S M W" ordering as fallback.
 
 Each function returns a list of floats (one per completion).
 """
 
-import math
 import os
 import re
 import sys
@@ -46,95 +41,57 @@ VALID_ACTIONS = {at.value.upper(): at for at in ActionType}
 DEFAULT_BELIEF = [0.5, 0.5, 0.5]
 
 
+_ANSWER_PATTERN = re.compile(
+    r"(\d)\s+(\d)\s+(\d)\s+(" + "|".join(at.value.upper() for at in ActionType) + r")\b",
+    re.IGNORECASE,
+)
+
+
 def extract_action_and_belief(text: str) -> tuple[ActionType | None, list[float], bool]:
-    """Parse 'ACTION_NAME S M W' format.
+    """Parse the agent's output and extract (action, belief, belief_provided).
+
+    Supports both formats:
+      - Plain answer line: "S M W ACTION_NAME"
+      - With CoT prefix:   "<reasoning>...</reasoning>\nS M W ACTION_NAME"
+
+    Strategy: search the entire response for the LAST occurrence of the
+    "<digit> <digit> <digit> <ACTION>" pattern. Taking the last match handles
+    cases where the model mentions an example mid-reasoning then commits to
+    a different answer at the end.
 
     Returns:
         (action, belief, belief_provided):
-            action: parsed ActionType or None
-            belief: 3-dim vector in [0, 1], DEFAULT_BELIEF if not parseable
-            belief_provided: True iff at least one belief digit was parsed from output
+            action: parsed ActionType or None if no valid action anywhere
+            belief: 3-dim vector in [0, 1], DEFAULT_BELIEF if no belief parsed
+            belief_provided: True iff a belief was extracted (full S M W ACTION pattern)
     """
     if not text:
         return None, list(DEFAULT_BELIEF), False
 
-    line = text.strip().split("\n")[0].strip()
-    if not line:
-        return None, list(DEFAULT_BELIEF), False
-
-    # Normalize: uppercase, comma → space
-    parts = line.upper().replace(",", " ").split()
-    if not parts:
-        return None, list(DEFAULT_BELIEF), False
-
-    # Find action and its index in parts
-    action: ActionType | None = None
-    action_idx = -1
-    for idx, p in enumerate(parts):
-        if p in VALID_ACTIONS:
-            action = VALID_ACTIONS[p]
-            action_idx = idx
-            break
-    if action is None:
-        for idx, p in enumerate(parts):
-            for name, at in VALID_ACTIONS.items():
-                if name in p:
-                    action = at
-                    action_idx = idx
-                    break
-            if action is not None:
-                break
-
-    # Iter 3: parse belief from BEFORE the action (belief-first format).
-    # Falls back to AFTER the action (legacy format) if no digits found before.
-    def _parse_digit(token: str) -> float | None:
-        token = token.strip().rstrip(".")
-        if not token:
-            return None
+    # Primary path: full belief+action pattern, take the LAST occurrence.
+    matches = list(_ANSWER_PATTERN.finditer(text))
+    if matches:
+        last = matches[-1]
+        s, m, w, action_name = last.groups()
         try:
-            if len(token) == 1 and token.isdigit():
-                return int(token) / 9.0
-            val = float(token)
-            if val > 1.0:
-                val = val / 9.0
-            return max(0.0, min(1.0, val))
-        except (ValueError, IndexError):
-            return None
-
-    belief: list[float] = []
-    belief_provided = False
-
-    if action_idx >= 0:
-        # Try belief-first: 3 digits BEFORE the action
-        if action_idx >= 3:
-            cand = [_parse_digit(parts[action_idx - 3 + i]) for i in range(3)]
-            if all(c is not None for c in cand):
-                belief = cand  # type: ignore[assignment]
-                belief_provided = True
-
-        # If belief-first didn't work, try legacy after-action format
-        if not belief_provided:
-            after_belief: list[float] = []
-            after_provided = False
-            for i in range(3):
-                j = action_idx + 1 + i
-                if j < len(parts):
-                    d = _parse_digit(parts[j])
-                    if d is not None:
-                        after_belief.append(d)
-                        after_provided = True
-                    else:
-                        after_belief.append(0.5)
-                else:
-                    after_belief.append(0.5)
-            if after_provided:
-                belief = after_belief
-                belief_provided = True
-
-    if not belief or len(belief) != 3:
-        belief = list(DEFAULT_BELIEF)
-
-    return action, belief, belief_provided
+            belief = [int(s) / 9.0, int(m) / 9.0, int(w) / 9.0]
+            action = ActionType(action_name.lower())
+            return action, belief, True
+        except (ValueError, KeyError):
+            pass  # fall through to action-only fallback
+
+    # Fallback: action-only output (no belief digits). Search for any valid
+    # ACTION_NAME token in the response and return that with default belief.
+    for line in text.strip().split("\n"):
+        for token in line.upper().replace(",", " ").split():
+            token = token.strip(".")
+            if token in VALID_ACTIONS:
+                return VALID_ACTIONS[token], list(DEFAULT_BELIEF), False
+            for name, at in VALID_ACTIONS.items():
+                if name in token:
+                    return at, list(DEFAULT_BELIEF), False
+
+    return None, list(DEFAULT_BELIEF), False
 
 
 def extract_action(text: str) -> ActionType | None:
@@ -168,10 +125,11 @@ def action_legal(completions, **kwargs) -> list[float]:
     """
     Layer 2: Is the parsed action one of the 10 valid actions?
 
-    All 10 actions are always legal in this env (no state-dependent validity).
-    Iter 2 fix: returns 0 for valid (was +0.5) so this layer becomes pure penalty
-    for malformed outputs. The +0.5 was a constant reward that contributed zero
-    to GRPO advantage and helped trigger mode collapse in iter 1.
+    All 10 actions are always legal in this env (no state-dependent validity),
+    so this layer is a pure penalty: 0 for any valid action, -1 for unparseable
+    output. Returning 0 (rather than a positive constant) is intentional —
+    a constant reward across all completions in a GRPO group contributes
+    exactly zero to the advantage and was a major contributor to mode collapse.
     """
     scores = []
     for completion in completions:
@@ -230,8 +188,9 @@ def env_reward(
                 ep_history = prompt_data.get("action_history", [])
                 ep_mode = prompt_data.get("profile_mode", "continuous")
             else:
-                # Iter 2 fix: mix index with prime to break deterministic seed clusters
-                # (avoids all completions in a position-class getting identical env_reward)
+                # Fallback seed mixes index with a prime to break deterministic
+                # seed clusters — without it, all completions in a position-class
+                # land on the same env_reward and contribute zero to GRPO advantage.
                 ep_seed = (i * 17) ^ 0xBEEF
                 ep_history = []
                 ep_mode = "continuous"
@@ -245,36 +204,38 @@ def env_reward(
             reward = obs.reward
             chosen = action_type.value
 
-            # Iter 4 fix (Issue 11): grader-aligned bias moved here from env.
-            # Per-step env reward stays pure (drives belief inference); the
-            # bias only shapes the GRPO-visible training reward.
+            # Grader-aligned bias (progress + connection deltas): shapes only
+            # the GRPO-visible training reward, not the env's per-step reward.
+            # Lives here rather than in env.step() so that the env's per-step
+            # reward (used by adaptation_score in the grader) stays pure.
             progress_delta = env._progress - pre_progress
             connection_delta = env._connection - pre_connection
             reward += 0.5 * progress_delta + 0.4 * connection_delta
 
-            # Iter 4 fix (Issue 3): SCALED-DOWN shaping. Iter 3 had
-            # -0.3/-0.4/+0.2 which dominated the ±0.5 env signal. Now
-            # roughly 1/3 of original magnitudes — nudges, not overrides.
+            # Diversity shaping: small nudges (~1/3 the magnitude of the env signal)
+            # so they don't dominate it. Three terms:
+            #   - repetition penalty (action ≥2× in last 3)
+            #   - low-entropy window penalty (last 6 actions ≤2 unique)
+            #   - new-action exploration bonus (until 6 distinct actions tried)
             if ep_history and len(ep_history) >= 2:
                 recent3 = ep_history[-3:]
                 if recent3.count(chosen) >= 2:
-                    reward -= 0.10  # was -0.3
+                    reward -= 0.10
 
             if ep_history and len(ep_history) >= 5:
                 last6 = ep_history[-5:] + [chosen]
                 if len(set(last6)) <= 2:
-                    reward -= 0.15  # was -0.4
+                    reward -= 0.15
 
             if ep_history is not None:
                 seen = set(ep_history)
                 if chosen not in seen and len(seen) < 6:
-                    reward += 0.07  # was +0.2
+                    reward += 0.07
 
-            # Iter 4 fix (Issue 10): BELIEF-ACTION COUPLING reward.
-            # Parse the agent's emitted belief and reward consistency between
-            # belief and action choice. Without this, the belief-first format
-            # only enforces consistency via causal attention (weak); now there's
-            # an explicit gradient signal.
+            # Belief-action coupling: rewards consistency between the agent's
+            # emitted belief and its chosen action. Without this term, the
+            # belief-first format only enforces consistency via causal attention
+            # (weak signal); this provides an explicit gradient.
             _, b, b_provided = extract_action_and_belief(response)
             if b_provided:
                 s_pref, m_pref, w_pref = b
@@ -308,7 +269,6 @@ def env_reward(
 
 def belief_accuracy(
     completions,
-    prompts=None,
     seed=None,
     step_index=None,
     action_history=None,
@@ -316,21 +276,17 @@ def belief_accuracy(
     **kwargs,
 ) -> list[float]:
     """
-    Layer 4: Belief-vector accuracy reward (META-LEARNING signal).
+    Layer 4: Belief-vector accuracy reward (the meta-learning signal).
 
-    ITER 4 FIX (Issue 4 from external review): The previous formula
-    `(1 - mae) - 0.5` gave constant emission "5 5 5" a free +0.336 reward
-    per step (× 3.0 weight = +1.0 per step = +28 per episode for ZERO learning).
-    This recreated the iter-1 collapse mechanism in disguise.
+    Reward = (1 − belief_mae) − constant_baseline_similarity, where the
+    constant baseline is what a "5 5 5" emission would score for THIS profile.
+    Subtracting the baseline matters: without it, a constant neutral emission
+    earns positive reward on every step (~+0.34 × the layer weight) for zero
+    actual learning, which silently re-creates the iter-1 mode-collapse pull.
+    With it: constant emission ≈ 0, better-than-baseline > 0, worse < 0.
 
-    New formula: subtract the per-profile baseline. The baseline is what a
-    constant 0.5 emission WOULD score for THIS profile. Now:
-      - Constant emission → reward ≈ 0 (no free reward)
-      - Better-than-baseline belief → positive
-      - Worse-than-baseline belief → negative
-
-    Plus iter 4 (Issue 9): no belief reward at step 0 (no information available
-    to commit a belief) — prevents pulling the policy toward a constant prior.
+    Belief reward is also skipped at step 0 — the agent has no information yet,
+    so penalizing belief-vs-target there just biases toward a constant prior.
     """
     scores = []
     for i, completion in enumerate(completions):
@@ -351,8 +307,9 @@ def belief_accuracy(
             scores.append(0.0)
             continue
 
-        # Iter 4 fix (Issue 9): step-0 commitment with no info biases toward
-        # constant prior. Skip belief reward at step 0.
+        # Step 0: agent has no information yet — committing a belief here
+        # would just pull the policy toward whatever constant prior the
+        # base model emits. Skip the reward for that step.
         if ep_step == 0:
             scores.append(0.0)
             continue
@@ -360,7 +317,6 @@ def belief_accuracy(
         try:
             env = _replay_env(ep_seed, ep_history, ep_mode)
             true_belief = env.get_belief_target()
-            # Iter 4 fix (Issue 4): subtract the constant-baseline reward
             mae = sum(abs(b - t) for b, t in zip(belief, true_belief)) / 3.0
             similarity = 1.0 - mae
             baseline_mae = sum(abs(0.5 - t) for t in true_belief) / 3.0
@@ -371,69 +327,3 @@ def belief_accuracy(
             scores.append(0.0)
 
     return scores
-
-
-def env_reward_simple(completions, prompts=None, **kwargs) -> list[float]:
-    """
-    State-aware heuristic reward (no env replay), used for fast smoke-training.
-    Identical to original simple reward — the format change is handled by extract_action.
-    """
-    scores = []
-    for i, completion in enumerate(completions):
-        response = completion[0]["content"] if isinstance(completion, list) else completion
-        action_type = extract_action(response)
-
-        if action_type is None:
-            scores.append(-2.0)
-            continue
-
-        prompt_text = ""
-        if prompts and i < len(prompts):
-            p = prompts[i]
-            if isinstance(p, list):
-                prompt_text = p[-1].get("content", "") if p else ""
-            elif isinstance(p, dict):
-                prompt_text = p.get("content", "")
-            else:
-                prompt_text = str(p)
-
-        meters = {}
-        for meter_name in ["Vitality", "Cognition", "Progress", "Serenity", "Connection"]:
-            match = re.search(rf"{meter_name}:\s*([\d.]+)", prompt_text)
-            if match:
-                meters[meter_name.lower()] = float(match.group(1))
-
-        is_morning = "Morning" in prompt_text
-        is_night = "Night" in prompt_text
-
-        score = 0.0
-        v = meters.get("vitality", 0.5)
-        s = meters.get("serenity", 0.5)
-        cn = meters.get("connection", 0.5)
-
-        if v < 0.2 and action_type in (ActionType.SLEEP, ActionType.EXERCISE):
-            score += 0.5
-        elif v < 0.2:
-            score -= 0.3
-
-        if s < 0.2 and action_type in (ActionType.MEDITATE, ActionType.ME_TIME):
-            score += 0.5
-        elif s < 0.2 and action_type != ActionType.SLEEP:
-            score -= 0.2
-
-        if cn < 0.2 and action_type in (ActionType.FAMILY_TIME, ActionType.SOCIALIZE):
-            score += 0.5
-        elif cn < 0.2:
-            score -= 0.1
-
-        if is_morning and action_type in (ActionType.DEEP_WORK, ActionType.LEARN):
-            score += 0.2
-        if is_night and action_type == ActionType.SLEEP:
-            score += 0.3
-
-        if action_type == ActionType.BINGE_WATCH:
-            score -= 0.3
-
-        scores.append(score)
-
-    return scores

training/sft_prime.py

@@ -0,0 +1,230 @@
+"""
+SFT prime: teach Qwen 2.5-3B the teacher's CoT-then-answer format.
+
+This is Stage 2 of Algorithm Distillation. We've already collected
+teacher trajectories (Stage 1). Here we fine-tune the student on the
+teacher's full responses — `<reasoning>...</reasoning>\nS M W ACTION_NAME` —
+so the student learns BOTH the format and the reasoning pattern that
+produced each answer.
+
+After this stage, the student should beat heuristic baselines on the
+v2 grader (which awards 0.20 for belief_accuracy). GRPO refinement is
+optional — only if the SFT'd model regresses on something.
+
+Usage (from rhythm_env root):
+    python training/sft_prime.py \
+        --teacher_jsonls data/teacher_30ep_validation.jsonl \
+                         data/teacher_indist_30_99.jsonl \
+                         data/teacher_ood_10000_10049.jsonl \
+        --output_dir outputs/rhythm-env-sft-primed \
+        --max_steps 600 \
+        --epochs 2
+
+Designed to run on HF Jobs with a10g-large flavor.
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+# The teacher's system prompt is the canonical contract — student must learn
+# to respond to this exact prompt. Imported from the teacher script for SSOT.
+from scripts.generate_teacher_trajectories import TEACHER_SYSTEM_PROMPT
+
+
+def load_teacher_dataset(jsonl_paths: list[str], drop_parse_fails: bool = True) -> list[dict]:
+    """Read teacher JSONL files and return list of {prompt, response} pairs.
+
+    Each input row is one step from one teacher episode. We turn it into a
+    chat-format SFT example: messages=[system, user] → completion=response.
+    Steps where the teacher's response failed to parse are dropped (we
+    don't want to teach the student bad outputs).
+    """
+    pairs: list[dict] = []
+    n_total = 0
+    n_dropped = 0
+    for path in jsonl_paths:
+        with open(path) as f:
+            for line in f:
+                row = json.loads(line)
+                n_total += 1
+                if drop_parse_fails and row.get("parse_failed"):
+                    n_dropped += 1
+                    continue
+                resp = row.get("teacher_response", "")
+                if not resp or not resp.strip():
+                    n_dropped += 1
+                    continue
+                pairs.append({
+                    "messages": [
+                        {"role": "system", "content": TEACHER_SYSTEM_PROMPT},
+                        {"role": "user", "content": row["user_prompt"]},
+                        {"role": "assistant", "content": resp},
+                    ],
+                })
+    print(f"Loaded {len(pairs)} SFT examples ({n_dropped}/{n_total} dropped: "
+          f"parse-failed or empty)")
+    return pairs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--teacher_jsonls", nargs="+", required=True,
+                        help="One or more teacher trajectory JSONL files")
+    parser.add_argument("--output_dir", type=str, default="outputs/rhythm-env-sft-primed")
+    parser.add_argument("--model_name", type=str, default="unsloth/Qwen2.5-3B-Instruct")
+    parser.add_argument("--epochs", type=int, default=2,
+                        help="SFT epochs over the dataset (2 is plenty for ~3000 examples)")
+    parser.add_argument("--max_steps", type=int, default=-1,
+                        help="Override epochs with a step count (-1 = use epochs)")
+    parser.add_argument("--lora_rank", type=int, default=16)
+    parser.add_argument("--learning_rate", type=float, default=2e-4)
+    parser.add_argument("--max_seq_length", type=int, default=2048,
+                        help="Must fit system + user + CoT response. ~600 user + ~120 CoT + ~10 ans + slack")
+    parser.add_argument("--per_device_batch_size", type=int, default=1)
+    parser.add_argument("--grad_accum", type=int, default=8,
+                        help="Effective batch size = per_device * grad_accum")
+    parser.add_argument("--warmup_ratio", type=float, default=0.1)
+    parser.add_argument("--save_method", type=str, default="merged_16bit",
+                        choices=["lora", "merged_16bit", "merged_4bit"])
+    args = parser.parse_args()
+
+    # ---- 1. Load + format the dataset ----
+    print("=" * 60)
+    print("Step 1: Loading teacher dataset")
+    print("=" * 60)
+    pairs = load_teacher_dataset(args.teacher_jsonls)
+    if not pairs:
+        sys.exit("ERROR: no SFT examples loaded — check JSONL paths")
+
+    from datasets import Dataset
+    raw_ds = Dataset.from_list(pairs)
+    print(f"Dataset size: {len(raw_ds)} examples")
+
+    # ---- 2. Load Qwen base via Unsloth ----
+    print("\n" + "=" * 60)
+    print(f"Step 2: Loading base model {args.model_name}")
+    print("=" * 60)
+    from unsloth import FastLanguageModel
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=args.model_name,
+        load_in_4bit=True,
+        max_seq_length=args.max_seq_length,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=args.lora_rank,
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        lora_alpha=args.lora_rank * 2,
+        use_gradient_checkpointing="unsloth",
+        random_state=3407,
+    )
+    print(f"LoRA rank {args.lora_rank}, alpha {args.lora_rank * 2}")
+
+    # ---- 3. Map to chat-template strings + tokenize ----
+    print("\n" + "=" * 60)
+    print("Step 3: Preparing dataset")
+    print("=" * 60)
+
+    def format_example(ex):
+        text = tokenizer.apply_chat_template(
+            ex["messages"],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        return {"text": text}
+
+    ds = raw_ds.map(format_example, remove_columns=raw_ds.column_names)
+    print("Sample formatted text (first 800 chars):")
+    print(ds[0]["text"][:800])
+    print("...")
+
+    # ---- 4. SFTTrainer ----
+    print("\n" + "=" * 60)
+    print("Step 4: Configuring SFTTrainer")
+    print("=" * 60)
+    from trl import SFTConfig, SFTTrainer
+
+    sft_kwargs = dict(
+        per_device_train_batch_size=args.per_device_batch_size,
+        gradient_accumulation_steps=args.grad_accum,
+        learning_rate=args.learning_rate,
+        warmup_ratio=args.warmup_ratio,
+        lr_scheduler_type="cosine",
+        optim="adamw_8bit",
+        weight_decay=0.001,
+        logging_steps=5,
+        save_strategy="no",
+        report_to="none",
+        output_dir=args.output_dir,
+        max_seq_length=args.max_seq_length,
+        dataset_text_field="text",
+        packing=False,
+    )
+    if args.max_steps > 0:
+        sft_kwargs["max_steps"] = args.max_steps
+    else:
+        sft_kwargs["num_train_epochs"] = args.epochs
+
+    sft_config = SFTConfig(**sft_kwargs)
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=ds,
+        args=sft_config,
+    )
+    print(f"Effective batch size: {args.per_device_batch_size * args.grad_accum}")
+    if args.max_steps > 0:
+        print(f"max_steps: {args.max_steps}")
+    else:
+        print(f"epochs: {args.epochs} → ~{len(ds) * args.epochs // (args.per_device_batch_size * args.grad_accum)} steps")
+
+    # ---- 5. Train ----
+    print("\n" + "=" * 60)
+    print("Step 5: Training")
+    print("=" * 60)
+    trainer.train()
+
+    # ---- 6. Save ----
+    print("\n" + "=" * 60)
+    print("Step 6: Saving model")
+    print("=" * 60)
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    if args.save_method == "lora":
+        model.save_pretrained(args.output_dir)
+        tokenizer.save_pretrained(args.output_dir)
+    else:
+        model.save_pretrained_merged(
+            args.output_dir,
+            tokenizer,
+            save_method=args.save_method,
+        )
+
+    # Save log_history for plot_from_log.py
+    log_path = os.path.join(args.output_dir, "log_history.json")
+    with open(log_path, "w") as f:
+        json.dump(trainer.state.log_history, f, indent=2)
+
+    # Save training config
+    config_path = os.path.join(args.output_dir, "training_config.json")
+    with open(config_path, "w") as f:
+        json.dump(vars(args), f, indent=2)
+
+    print(f"\nSaved SFT-primed model to: {args.output_dir}")
+    print(f"Log history:               {log_path}")
+    print(f"Training config:           {config_path}")
+    print()
+    print("Next: python training/inference_eval.py --model_path " + args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

training/train.py

@@ -13,8 +13,8 @@ Usage (Colab T4):
     !pip install unsloth transformers trl datasets
     !python training/train.py --max_steps 1500
 
-Usage (local — for setup checks; needs CUDA for actual training):
-    python training/train.py --max_steps 50 --use_simple_reward
+Setup-check (no GPU): run the smoke tests instead of starting a real run:
+    python -m pytest tests/test_pipeline_smoke.py -q
 """
 
 import argparse
@@ -45,12 +45,7 @@ def main():
                         help="Fraction of dataset with profile hint visible. Default 0.0 (no hints) "
                              "to eliminate train-eval distribution mismatch. Set >0 only if you ALSO "
                              "show hints during eval.")
-    parser.add_argument("--profile_mode", type=str, default="continuous",
-                        choices=["continuous", "discrete"],
-                        help="continuous = sampled per-episode (meta-RL); discrete = 3 hardcoded profiles")
     parser.add_argument("--output_dir", type=str, default="outputs/rhythmenv_meta_trained")
-    parser.add_argument("--use_simple_reward", action="store_true",
-                        help="Use heuristic reward instead of env-replay (smoke testing)")
     parser.add_argument("--report_to", type=str, default="none")
     args = parser.parse_args()
 
@@ -68,7 +63,6 @@ def main():
         num_episodes=args.num_episodes,
         strategy="mixed",
         max_samples=args.max_samples,
-        profile_mode=args.profile_mode,
         hint_fraction=args.hint_fraction,
     )
 
@@ -123,16 +117,10 @@ def main():
     print("Step 3: Setting up reward functions")
     print("=" * 60)
 
-    from reward_functions import (
-        format_valid, action_legal, env_reward, env_reward_simple, belief_accuracy
-    )
+    from reward_functions import format_valid, action_legal, env_reward, belief_accuracy
 
-    if args.use_simple_reward:
-        reward_funcs = [format_valid, action_legal, env_reward_simple, belief_accuracy]
-        print("Using: format_valid + action_legal + env_reward_simple + belief_accuracy")
-    else:
-        reward_funcs = [format_valid, action_legal, env_reward, belief_accuracy]
-        print("Using: format_valid + action_legal + env_reward + belief_accuracy")
+    reward_funcs = [format_valid, action_legal, env_reward, belief_accuracy]
+    print("Using: format_valid + action_legal + env_reward + belief_accuracy")
 
     # ---------------------------------------------------------------
     # 4. GRPO trainer config
@@ -146,11 +134,11 @@ def main():
     max_prompt_length = 600          # history + hint room
     max_completion_length = 32       # bumped from 20 to prevent silent truncation of belief digits
 
-    # reward_weights: aggressively suppress saturated format/legal layers, amplify
-    # the variable signals (env + belief). Iter 1 collapsed because format/legal at
-    # +1.0/+0.5 contributed zero to GRPO advantage; iter 2 makes belief the dominant
-    # learning signal.
-    # Order MUST match reward_funcs in main(): format_valid, action_legal, env_reward, belief_accuracy
+    # reward_weights: suppress the format/action_legal layers (small, low-variance
+    # signals — too constant across a GRPO group to contribute meaningful advantage)
+    # and amplify the variable signals env_reward and belief_accuracy. belief_accuracy
+    # at 3.0 is the dominant learning signal.
+    # Order MUST match reward_funcs above: format_valid, action_legal, env_reward, belief_accuracy
     reward_weights = [0.05, 0.05, 1.5, 3.0]
 
     training_args_kwargs = dict(