Rewrite notebook: per-piece GRPO training on Qwen 3B

Complete rewrite of training approach:
- Model sees board before EVERY piece (not 100 blind actions)
- Generates up to 20 action tokens per piece, stops when piece locks
- Custom REINFORCE/GRPO loop (no GRPOTrainer dependency)
- Two-phase: rollout without grad, then recompute log_probs with grad
- 8 games per iteration (same seed, GRPO-style comparison)
- 100 iterations with different seeds

Model changes:
- Qwen2.5-3B-Instruct (was 7B)
- LoRA r=16 (was r=32)
- Target GPU: L4 (was A100)

Reward structure:
- Engine per-step rewards summed across entire game
- L/R penalty: -0.1 per move
- Piece placed bonus: +1.0
- No-place penalty: -10.0 (forced drop after 20 tokens)
- Line clear bonus: +100 per line

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

docs/plans/2026-03-09-per-piece-training-design.md

@@ -0,0 +1,172 @@
+# Per-Piece Tetris Training Design
+
+## Overview
+
+Redesign training so the model sees the board before EVERY piece placement,
+instead of outputting 100 blind actions from a single prompt.
+
+## Key Parameters
+
+- **Model**: Qwen2.5-3B-Instruct + LoRA
+- **GPU**: L4
+- **Games per iteration**: 8 (same seed, GRPO-style comparison)
+- **Max steps per game**: 200 (total actions across all pieces)
+- **Max tokens per piece**: 20 (if piece not placed -> forced drop + penalty)
+- **Training iterations**: 100 (each with a new seed)
+- **No history**: each model call = fresh prompt with current board only
+
+## One Model Call (one piece)
+
+1. Build prompt: system message + current board + current piece + next piece
+2. Model generates up to 20 tokens (L/R/C/W/D/S)
+3. Play actions one by one on the engine
+4. Stop when piece locks (new piece spawns) or 20 tokens exhausted
+5. If 20 tokens used and piece NOT placed: forced drop + penalty (-10)
+
+### How to detect piece lock
+
+After each `env.step()`, check if `current_piece` changed from the one
+shown in the prompt. If it changed -> piece was placed, stop processing tokens.
+
+## One Game (one playthrough)
+
+```
+seed = iteration_seed
+env.reset(seed)
+total_reward = 0
+all_log_probs = []
+steps = 0
+
+while not game_over and steps < 200:
+    board_state = env.get_state()
+    current_piece = board_state['current_piece']
+
+    prompt = build_prompt(board_state)  # fresh each time, no history
+    tokens, log_probs = model_generate(prompt, max_tokens=20)
+
+    piece_placed = False
+    for token in tokens:
+        action = token_to_action(token)
+        result = env.step(action)
+        total_reward += result['reward']
+        total_reward -= 0.1 if action in ('left', 'right') else 0
+        all_log_probs.append(log_probs[token])
+        steps += 1
+
+        if result['current_piece'] != current_piece:
+            # Piece was placed, new piece spawned
+            piece_placed = True
+            total_reward += 1.0  # bonus for placing piece
+            break
+
+        if result['done']:
+            game_over = True
+            break
+
+    if not piece_placed and not game_over:
+        # Piece not placed in 20 tokens -> force drop + penalty
+        env.step('drop')
+        total_reward -= 10.0
+        steps += 1
+```
+
+## Reward Structure
+
+Per-step rewards from game engine (summed across ALL steps):
+- Step penalty: -1 per step
+- Line clears: +100/+300/+700/+1500 (1/2/3/4 lines)
+- Height penalty: -2 * max_height (per step)
+- Hole penalty: -5 * holes (per step)
+- Game over: -500
+
+Additional rewards added by training loop:
+- L or R move: -0.1 (discourages aimless shuffling)
+- Piece placed: +1.0 (encourages completing placements)
+- Piece NOT placed in 20 tokens: -10.0 (forces learning to drop)
+- Line clear bonus: +100 per line (1 line=+100, 2=+200, 3=+300, 4=+400)
+
+Total game reward = sum of all the above across entire game.
+
+## One Training Iteration
+
+```
+seed = random_seed_for_this_iteration
+
+# Play 8 games with same seed
+rewards = []
+all_game_log_probs = []
+
+for game_idx in range(8):
+    reward, log_probs_sum = play_one_game(model, seed)
+    rewards.append(reward)
+    all_game_log_probs.append(log_probs_sum)
+
+# GRPO-style advantage
+rewards = torch.tensor(rewards)
+advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8)
+
+# Policy gradient loss
+loss = 0
+for i in range(8):
+    loss -= all_game_log_probs[i] * advantages[i]
+loss = loss / 8
+
+loss.backward()
+optimizer.step()
+optimizer.zero_grad()
+```
+
+## Full Training
+
+```
+for iteration in range(100):
+    seed = iteration  # deterministic but different each time
+    train_one_iteration(model, seed)
+
+    if iteration % 10 == 0:
+        print(f"Iter {iteration}: avg_reward={mean}, std={std}")
+```
+
+## Action Token Mapping
+
+Pre-compute token IDs for action characters:
+```python
+ACTION_TOKENS = {
+    tokenizer.encode('L', add_special_tokens=False)[0]: 'left',
+    tokenizer.encode('R', add_special_tokens=False)[0]: 'right',
+    tokenizer.encode('C', add_special_tokens=False)[0]: 'rotate_cw',
+    tokenizer.encode('W', add_special_tokens=False)[0]: 'rotate_ccw',
+    tokenizer.encode('D', add_special_tokens=False)[0]: 'drop',
+    tokenizer.encode('S', add_special_tokens=False)[0]: 'down',
+}
+```
+
+During generation: mask logits to only allow these 6 tokens.
+Sample from softmax over 6 logits -> get action + log_prob.
+
+## Notebook Cell Structure
+
+1. Install deps (peft, trl, accelerate, etc.)
+2. Load Qwen2.5-3B-Instruct + LoRA
+3. Download game_engine.py, define prompt builder
+4. **Demo: untrained model plays one game** (show board after each piece)
+5. Define `play_one_game()` and `train_one_iteration()`
+6. Training loop (100 iterations)
+7. Plot reward curve
+8. **Demo: trained model plays one game** (compare with untrained)
+9. Push model to HF Hub
+
+## Time Estimate
+
+- 3B model forward pass: ~0.05s per token on T4
+- Per piece: ~20 tokens * 0.05s = ~1s (forward passes)
+- Per game: ~30-50 pieces * 1s = ~40s
+- Per iteration: 8 games (sequential) = ~320s OR batched = ~40-80s
+- 100 iterations: ~70-130 min
+
+Fits in a Colab T4 session (4h limit).
+
+Note: batching 8 games in parallel requires handling variable-length episodes.
+Simpler approach: run 8 games sequentially (~5 min/iteration, ~8h total).
+Better approach: batch the forward passes across 8 games at each "step",
+masking out finished games.

docs/plans/2026-03-09-per-piece-training-plan.md

@@ -0,0 +1,580 @@
+# Per-Piece Tetris GRPO Training — Implementation Plan
+
+> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
+
+**Goal:** Rewrite the Colab notebook so the model sees the board before every piece, plays actions until piece locks, and learns via GRPO-style policy gradient over full games.
+
+**Architecture:** Custom REINFORCE/GRPO loop (no GRPOTrainer). Two-phase per iteration: (1) rollout 8 games without grad, storing (prompt, actions) pairs; (2) recompute log_probs with grad, multiply by advantage, backward. This keeps memory low — only one piece's activations in memory at a time.
+
+**Tech Stack:** transformers, peft (LoRA), torch, game_engine.py (local Tetris)
+
+---
+
+### Task 1: Cell 1 — Install Dependencies
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 1
+
+**Step 1: Write cell content**
+
+```python
+# Cell 1: Install dependencies
+!pip install peft accelerate -q
+```
+
+Note: we no longer need `trl` (no GRPOTrainer), `openenv-core`, or `datasets`.
+
+**Step 2: Run cell, verify no errors**
+
+**Step 3: Commit**
+
+```bash
+git add tetris_training.ipynb
+git commit -m "cell 1: minimal deps for custom training loop"
+```
+
+---
+
+### Task 2: Cell 2 — Load Model + LoRA
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 2
+
+**Step 1: Write cell content**
+
+```python
+# Cell 2: Load Qwen2.5-3B-Instruct + LoRA
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+
+model_name = "Qwen/Qwen2.5-3B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+lora_config = LoraConfig(
+    r=16,
+    lora_alpha=16,
+    lora_dropout=0,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                    "gate_proj", "up_proj", "down_proj"],
+    task_type="CAUSAL_LM",
+)
+
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+```
+
+Note: 3B model, r=16 (smaller than 7B's r=32), bf16.
+
+**Step 2: Run cell, verify output shows trainable params**
+
+**Step 3: Commit**
+
+---
+
+### Task 3: Cell 3 — Game Engine + Constants + Prompt Builder
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 3 (replaces old TetrisClient cell)
+
+**Step 1: Write cell content**
+
+```python
+# Cell 3: Game engine + constants + prompt builder
+import random
+import torch.nn.functional as F
+
+# Download game engine
+!wget -q -O game_engine.py https://raw.githubusercontent.com/OutOfMystic/tetris-openenv/main/src/tetris_env/server/game_engine.py
+from game_engine import TetrisEnv
+
+# === Constants ===
+MAX_ACTIONS_PER_PIECE = 20
+MAX_STEPS_PER_GAME = 200
+GAMES_PER_ITER = 8
+NUM_ITERATIONS = 100
+TEMPERATURE = 0.7
+
+# Training reward modifiers (on top of engine rewards)
+LR_PENALTY = -0.1          # per L/R move
+PIECE_PLACED_BONUS = 1.0   # per piece successfully placed
+NO_PLACE_PENALTY = -10.0   # if 20 tokens exhausted without placing
+LINE_CLEAR_BONUS = 100.0   # per line cleared (1=+100, 2=+200, 3=+300, 4=+400)
+
+# Action mapping
+ACTION_CHARS = ['L', 'R', 'C', 'W', 'D', 'S']
+ACTION_TO_ENGINE = {
+    'L': 'left', 'R': 'right', 'C': 'rotate_cw',
+    'W': 'rotate_ccw', 'D': 'drop', 'S': 'down'
+}
+
+# Pre-compute token IDs for action characters
+ACTION_TOKEN_IDS = []
+for ch in ACTION_CHARS:
+    ids = tokenizer.encode(ch, add_special_tokens=False)
+    ACTION_TOKEN_IDS.append(ids[0])
+    print(f"  '{ch}' -> token_id {ids[0]}")
+ACTION_TOKEN_IDS = torch.tensor(ACTION_TOKEN_IDS, device=model.device)
+
+# Token ID -> action char lookup
+TOKEN_TO_CHAR = {ACTION_TOKEN_IDS[i].item(): ACTION_CHARS[i] for i in range(6)}
+
+SYSTEM_PROMPT = """You are a Tetris AI. You see the board and current piece.
+Output actions as single letters: L=left R=right C=rotate_cw W=rotate_ccw D=drop S=down
+Place the piece to fill complete rows. Drop when positioned."""
+
+def build_prompt(result):
+    return f"""Board:
+{result['board']}
+
+Piece: {result['current_piece']} Next: {result['next_piece']}
+Score: {result['score']} Lines: {result['total_lines']} Height: {result['max_height']} Holes: {result['holes']}
+
+Your actions:"""
+
+print("Game engine loaded. Action tokens mapped.")
+```
+
+**Step 2: Run cell, verify 6 token mappings printed**
+
+**Step 3: Commit**
+
+---
+
+### Task 4: Cell 4 — Core Functions (play_one_game + train_one_iteration)
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 4
+
+This is the most critical cell. Two phases per iteration:
+- **Rollout** (no grad): play 8 games, store (prompt_ids, action_token_ids) per piece
+- **Update** (with grad): recompute log_probs, multiply by advantage, backward per piece
+
+**Step 1: Write cell content**
+
+```python
+# Cell 4: Core training functions
+
+def play_one_game(model, tokenizer, seed, temperature=TEMPERATURE):
+    """
+    Play a full Tetris game. Model sees board before each piece,
+    generates up to 20 action tokens per piece.
+    Returns reward and piece data for gradient computation.
+    """
+    env = TetrisEnv(seed=seed)
+    env.reset(seed=seed)
+
+    # Initial random offset (same as old prompt generation)
+    rng = random.Random(seed)
+    moves = rng.randint(0, 4)
+    direction = rng.choice(["left", "right"])
+    for _ in range(moves):
+        if env.done:
+            break
+        env.step(direction)
+
+    total_reward = 0.0
+    total_steps = 0
+    pieces_data = []
+
+    while not env.done and total_steps < MAX_STEPS_PER_GAME:
+        current_piece_name = env.current_piece_name
+        lines_before = env.total_lines
+
+        # Build fresh prompt for this piece
+        result = env._make_result(0)
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": build_prompt(result)},
+        ]
+        prompt_ids = tokenizer.apply_chat_template(
+            messages, return_tensors="pt", add_generation_prompt=True
+        ).to(model.device)
+
+        # Autoregressive generation with KV cache (no grad)
+        input_ids = prompt_ids
+        action_ids_list = []
+        piece_placed = False
+        past_kv = None
+
+        with torch.no_grad():
+            for _ in range(MAX_ACTIONS_PER_PIECE):
+                if past_kv is None:
+                    out = model(input_ids, use_cache=True)
+                    past_kv = out.past_key_values
+                else:
+                    out = model(input_ids[:, -1:], past_key_values=past_kv, use_cache=True)
+                    past_kv = out.past_key_values
+
+                logits = out.logits[:, -1, :]  # [1, vocab]
+
+                # Mask to only 6 action tokens
+                masked = torch.full_like(logits, float('-inf'))
+                masked[0, ACTION_TOKEN_IDS] = logits[0, ACTION_TOKEN_IDS]
+                probs = F.softmax(masked / temperature, dim=-1)
+
+                # Sample
+                token_id = torch.multinomial(probs, 1).item()
+                action_ids_list.append(token_id)
+
+                # Execute in engine
+                action_char = TOKEN_TO_CHAR[token_id]
+                action_name = ACTION_TO_ENGINE[action_char]
+                step_result = env.step(action_name)
+
+                total_reward += step_result['reward']
+                total_steps += 1
+
+                if action_name in ('left', 'right'):
+                    total_reward += LR_PENALTY
+
+                # Check if piece placed (new piece spawned)
+                if env.current_piece_name != current_piece_name:
+                    piece_placed = True
+                    total_reward += PIECE_PLACED_BONUS
+                    lines_cleared = env.total_lines - lines_before
+                    if lines_cleared > 0:
+                        total_reward += lines_cleared * LINE_CLEAR_BONUS
+                    break
+
+                if env.done:
+                    break
+
+                # Append token for next autoregressive step
+                input_ids = torch.cat([
+                    input_ids,
+                    torch.tensor([[token_id]], device=model.device)
+                ], dim=-1)
+
+        # Force drop if piece not placed
+        if not piece_placed and not env.done:
+            env.step('drop')
+            total_reward += NO_PLACE_PENALTY
+            total_steps += 1
+            lines_cleared = env.total_lines - lines_before
+            if lines_cleared > 0:
+                total_reward += lines_cleared * LINE_CLEAR_BONUS
+
+        # Store for gradient computation
+        if action_ids_list:
+            pieces_data.append({
+                'prompt_ids': prompt_ids.cpu(),
+                'action_ids': torch.tensor(action_ids_list, dtype=torch.long),
+            })
+
+    return {
+        'reward': total_reward,
+        'pieces': pieces_data,
+        'total_steps': total_steps,
+        'total_lines': env.total_lines,
+        'pieces_placed': len(pieces_data),
+    }
+
+
+def train_one_iteration(model, optimizer, seed, temperature=TEMPERATURE):
+    """
+    One GRPO iteration:
+    1. Play 8 games (same seed) without grad
+    2. Compute advantages
+    3. Recompute log_probs with grad, apply policy gradient
+    """
+    # Phase 1: Rollout
+    games = []
+    for _ in range(GAMES_PER_ITER):
+        game = play_one_game(model, tokenizer, seed, temperature)
+        games.append(game)
+
+    rewards = torch.tensor([g['reward'] for g in games], dtype=torch.float32)
+    mean_r = rewards.mean().item()
+    std_r = rewards.std().item()
+
+    # Phase 2: Advantages (GRPO-style)
+    if std_r < 1e-8:
+        # All games got same reward — no learning signal
+        return {'mean_reward': mean_r, 'std_reward': 0.0, 'loss': 0.0,
+                'avg_steps': sum(g['total_steps'] for g in games) / GAMES_PER_ITER,
+                'avg_lines': sum(g['total_lines'] for g in games) / GAMES_PER_ITER}
+
+    advantages = ((rewards - rewards.mean()) / (rewards.std() + 1e-8)).tolist()
+
+    # Phase 3: Update — recompute log_probs with grad
+    optimizer.zero_grad()
+    total_pieces = sum(len(g['pieces']) for g in games)
+    total_loss = 0.0
+
+    for game_idx, game in enumerate(games):
+        adv = advantages[game_idx]
+        for piece in game['pieces']:
+            prompt = piece['prompt_ids'].to(model.device)
+            actions = piece['action_ids'].to(model.device)
+            if len(actions) == 0:
+                continue
+
+            # Teacher-forced forward pass: prompt + actions
+            full_input = torch.cat([prompt.squeeze(0), actions]).unsqueeze(0)
+            logits = model(full_input).logits
+
+            # Log_probs at positions where actions were generated
+            P = prompt.shape[-1]
+            action_logits = logits[0, P-1 : P-1+len(actions), :]
+
+            # Mask to action tokens, apply temperature
+            masked = torch.full_like(action_logits, float('-inf'))
+            masked[:, ACTION_TOKEN_IDS] = action_logits[:, ACTION_TOKEN_IDS]
+            log_probs = F.log_softmax(masked / temperature, dim=-1)
+
+            selected = log_probs.gather(1, actions.unsqueeze(1)).squeeze(1)
+            piece_loss = -(selected.sum() * adv) / total_pieces
+
+            piece_loss.backward()
+            total_loss += piece_loss.item()
+
+    # Gradient clipping
+    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+    optimizer.step()
+
+    return {
+        'mean_reward': mean_r,
+        'std_reward': std_r,
+        'loss': total_loss,
+        'avg_steps': sum(g['total_steps'] for g in games) / GAMES_PER_ITER,
+        'avg_lines': sum(g['total_lines'] for g in games) / GAMES_PER_ITER,
+        'avg_pieces': sum(g['pieces_placed'] for g in games) / GAMES_PER_ITER,
+    }
+
+print("Training functions defined.")
+```
+
+**Step 2: Run cell, verify "Training functions defined." printed**
+
+**Step 3: Quick smoke test**
+
+```python
+# Smoke test: play 1 game
+test_game = play_one_game(model, tokenizer, seed=0)
+print(f"Reward: {test_game['reward']:.1f}, Steps: {test_game['total_steps']}, "
+      f"Pieces: {test_game['pieces_placed']}, Lines: {test_game['total_lines']}")
+```
+
+**Step 4: Commit**
+
+---
+
+### Task 5: Cell 5 — Demo Untrained Model
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 5
+
+**Step 1: Write cell content**
+
+```python
+# Cell 5: Demo — UNTRAINED model plays one game
+print("=== UNTRAINED MODEL ===\n")
+
+game = play_one_game(model, tokenizer, seed=42)
+
+print(f"Total steps: {game['total_steps']}")
+print(f"Pieces placed: {game['pieces_placed']}")
+print(f"Lines cleared: {game['total_lines']}")
+print(f"Game reward: {game['reward']:+.1f}")
+
+# Show final board
+env = TetrisEnv(seed=42)
+env.reset(seed=42)
+rng = random.Random(42)
+moves = rng.randint(0, 4)
+direction = rng.choice(["left", "right"])
+for _ in range(moves):
+    env.step(direction)
+
+# Replay the actions to get final board
+for piece in game['pieces']:
+    for token_id in piece['action_ids'].tolist():
+        action_name = ACTION_TO_ENGINE[TOKEN_TO_CHAR[token_id]]
+        if not env.done:
+            env.step(action_name)
+    if not env.done:
+        # Check if piece was placed; if not, force drop
+        pass  # engine handles lock internally
+
+print(f"\nFinal board:")
+print(env.board_to_text())
+
+untrained_reward = game['reward']
+```
+
+**Step 2: Run cell, verify board prints**
+
+**Step 3: Commit**
+
+---
+
+### Task 6: Cell 6 — Training Loop
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 6
+
+**Step 1: Write cell content**
+
+```python
+# Cell 6: Training loop — 100 iterations
+optimizer = torch.optim.AdamW(
+    [p for p in model.parameters() if p.requires_grad],
+    lr=5e-6,
+    weight_decay=0.01,
+)
+
+history = []
+
+print("Starting per-piece GRPO training...")
+print(f"Config: {GAMES_PER_ITER} games/iter, max {MAX_STEPS_PER_GAME} steps, "
+      f"max {MAX_ACTIONS_PER_PIECE} tokens/piece, {NUM_ITERATIONS} iterations\n")
+
+for iteration in range(NUM_ITERATIONS):
+    stats = train_one_iteration(model, optimizer, seed=iteration)
+    history.append(stats)
+
+    if iteration % 5 == 0 or iteration == NUM_ITERATIONS - 1:
+        print(f"[Iter {iteration:3d}] "
+              f"reward={stats['mean_reward']:+8.1f} "
+              f"std={stats['std_reward']:6.1f} "
+              f"loss={stats['loss']:7.3f} "
+              f"steps={stats['avg_steps']:5.1f} "
+              f"lines={stats['avg_lines']:4.1f} "
+              f"pieces={stats['avg_pieces']:4.1f}")
+
+print("\nTraining complete!")
+```
+
+**Step 2: Run cell, verify reward logs appear**
+
+**Step 3: Commit**
+
+---
+
+### Task 7: Cell 7 — Plot Reward Curve
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 7
+
+**Step 1: Write cell content**
+
+```python
+# Cell 7: Plot reward curve
+import matplotlib.pyplot as plt
+
+fig, axes = plt.subplots(1, 3, figsize=(18, 5))
+
+iters = range(len(history))
+
+axes[0].plot(iters, [h['mean_reward'] for h in history])
+axes[0].set_title('Mean Reward per Iteration')
+axes[0].set_xlabel('Iteration')
+axes[0].set_ylabel('Reward')
+
+axes[1].plot(iters, [h['avg_lines'] for h in history])
+axes[1].set_title('Avg Lines Cleared')
+axes[1].set_xlabel('Iteration')
+
+axes[2].plot(iters, [h['loss'] for h in history])
+axes[2].set_title('Policy Gradient Loss')
+axes[2].set_xlabel('Iteration')
+
+plt.tight_layout()
+plt.savefig('reward_curve.png', dpi=150)
+plt.show()
+```
+
+**Step 2: Run, verify plots**
+
+**Step 3: Commit**
+
+---
+
+### Task 8: Cell 8 — Demo Trained Model
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 8
+
+**Step 1: Write cell content**
+
+```python
+# Cell 8: Demo — TRAINED model plays 3 games
+print("=== TRAINED MODEL ===\n")
+
+trained_rewards = []
+for seed in [42, 123, 7]:
+    game = play_one_game(model, tokenizer, seed=seed)
+    print(f"Seed {seed}: reward={game['reward']:+.1f}, "
+          f"steps={game['total_steps']}, lines={game['total_lines']}, "
+          f"pieces={game['pieces_placed']}")
+    trained_rewards.append(game['reward'])
+
+avg_trained = sum(trained_rewards) / len(trained_rewards)
+print(f"\n{'='*50}")
+print(f"UNTRAINED reward (seed=42): {untrained_reward:+.1f}")
+print(f"TRAINED avg reward (3 games): {avg_trained:+.1f}")
+print(f"Improvement: {avg_trained - untrained_reward:+.1f}")
+print('='*50)
+```
+
+**Step 2: Run, verify comparison**
+
+**Step 3: Commit**
+
+---
+
+### Task 9: Cell 9 — Push to HF Hub
+
+**Files:**
+- Modify: `tetris_training.ipynb` Cell 9
+
+**Step 1: Write cell content**
+
+```python
+# Cell 9: Push trained model to HF Hub
+model.push_to_hub("VortexedSquirrel/tetris-agent-grpo")
+tokenizer.push_to_hub("VortexedSquirrel/tetris-agent-grpo")
+print("Model pushed to https://huggingface.co/VortexedSquirrel/tetris-agent-grpo")
+```
+
+**Step 2: Commit + push**
+
+```bash
+git add tetris_training.ipynb
+git commit -m "rewrite notebook: per-piece GRPO training on 3B"
+git push origin main
+```
+
+---
+
+## Key Implementation Details
+
+### Piece lock detection
+After `env.step()`, check `env.current_piece_name != current_piece_name`.
+The engine calls `_spawn_next()` when a piece locks, changing the current piece.
+
+### Two-phase gradient computation
+- **Rollout phase**: `torch.no_grad()`, KV cache for fast autoregressive generation
+- **Update phase**: teacher-forced forward pass (prompt + actions in one call),
+  `piece_loss.backward()` after each piece to keep memory low
+
+### Memory-safe gradient accumulation
+Each piece does its own `.backward()` adding to accumulated gradients.
+Only one piece's computation graph in memory at a time.
+`optimizer.step()` called once after all 8 games processed.
+
+### Time estimate (L4, 3B model)
+- Rollout: 8 games × ~200 steps × ~0.02s/token ≈ 32s
+- Update: 8 games × ~30 pieces × ~0.1s/piece ≈ 24s
+- Total per iteration: ~56s
+- 100 iterations: ~93 min ≈ 1.5 hours

tetris_training.ipynb

@@ -5,11 +5,11 @@
    "metadata": {
     "id": "jYiMCaOSFG8J"
    },
-   "source": "# Tetris OpenEnv — GRPO Training\n\nTrain an LLM agent to play Tetris using GRPO (Group Relative Policy Optimization).\n\n**Environment**: Tetris on HF Spaces via OpenEnv 0.2.1\n**Model**: Qwen2.5-7B-Instruct + LoRA (standard HF stack)\n**Training**: GRPO via TRL — single-prompt, 100-action blind play\n**Runtime**: A100 GPU (Colab Pro)"
+   "source": "# Tetris OpenEnv — Per-Piece GRPO Training\n\nTrain an LLM to play Tetris: model sees the board before **every piece**, outputs actions until piece locks, learns via GRPO-style policy gradient.\n\n**Environment**: Local Tetris engine (game_engine.py from OpenEnv)\n**Model**: Qwen2.5-3B-Instruct + LoRA (r=16)\n**Training**: Custom REINFORCE/GRPO loop — 8 games/iteration, 100 iterations\n**Runtime**: L4 GPU (Colab)"
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -17,35 +17,12 @@
     "id": "PHNUG6nYFG8L",
     "outputId": "e913c80a-f0ec-4231-a35c-f01106a333a1"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "\u001b[33mWARNING: Skipping unsloth as it is not installed.\u001b[0m\u001b[33m\n",
-      "\u001b[0m\u001b[33mWARNING: Skipping unsloth-zoo as it is not installed.\u001b[0m\u001b[33m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m528.8/528.8 kB\u001b[0m \u001b[31m10.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m121.9/121.9 kB\u001b[0m \u001b[31m15.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m633.7/633.7 kB\u001b[0m \u001b[31m33.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m251.7/251.7 kB\u001b[0m \u001b[31m20.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m201.9/201.9 kB\u001b[0m \u001b[31m23.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m96.4/96.4 kB\u001b[0m \u001b[31m11.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m152.3/152.3 kB\u001b[0m \u001b[31m17.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m80.2/80.2 kB\u001b[0m \u001b[31m9.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m331.1/331.1 kB\u001b[0m \u001b[31m34.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[?25h"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 1: Install dependencies (remove Unsloth, use standard HF stack)\n",
-    "!pip uninstall unsloth unsloth-zoo -y -q\n",
-    "!pip install peft trl openenv-core datasets accelerate -q\n"
-   ]
+   "outputs": [],
+   "source": "# Cell 1: Install dependencies\n!pip install peft accelerate -q"
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
@@ -166,208 +143,12 @@
     "id": "4CEG1_JMFG8L",
     "outputId": "fe327915-556f-409f-c75c-fc2f74a7e5f6"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stderr",
-     "text": [
-      "/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n",
-      "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
-      "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
-      "You will be able to reuse this secret in all of your notebooks.\n",
-      "Please note that authentication is recommended but still optional to access public models or datasets.\n",
-      "  warnings.warn(\n"
-     ]
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "config.json:   0%|          | 0.00/663 [00:00<?, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "6285b80f4b53409094ba1877193efba5"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "tokenizer_config.json: 0.00B [00:00, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "c34a5bdf556842fa9faaab6bac674e4b"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "vocab.json: 0.00B [00:00, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "ed82c562c8824a1ebf459cdcbad610c9"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "merges.txt: 0.00B [00:00, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "fcd3965e71074108a5e932cbb97c40bc"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "tokenizer.json: 0.00B [00:00, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "6c8415ffb68445619f115c60e32c616b"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "stream",
-     "name": "stderr",
-     "text": [
-      "`torch_dtype` is deprecated! Use `dtype` instead!\n"
-     ]
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "model.safetensors.index.json: 0.00B [00:00, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "f7918c9c71a34e178a45e889e33cec93"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "Downloading (incomplete total...): 0.00B [00:00, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "8c7e3c13b85d4013bd571019d6efff5f"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "Fetching 4 files:   0%|          | 0/4 [00:00<?, ?it/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "ff3dd07d11e04ca8b5bef7c77d95c7ff"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "Loading weights:   0%|          | 0/339 [00:00<?, ?it/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "cfe6c9c2bc5d460d846f2af15f9d7ed6"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "generation_config.json:   0%|          | 0.00/243 [00:00<?, ?B/s]"
-      ],
-      "application/vnd.jupyter.widget-view+json": {
-       "version_major": 2,
-       "version_minor": 0,
-       "model_id": "318e1f3e9c4f4c5992a1507cd9de6d9a"
-      }
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "trainable params: 80,740,352 || all params: 7,696,356,864 || trainable%: 1.0491\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 2: Load model (standard HF stack — no Unsloth bug)\n",
-    "import torch\n",
-    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
-    "from peft import LoraConfig, get_peft_model\n",
-    "\n",
-    "model_name = \"Qwen/Qwen2.5-7B-Instruct\"\n",
-    "\n",
-    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
-    "if tokenizer.pad_token is None:\n",
-    "    tokenizer.pad_token = tokenizer.eos_token\n",
-    "\n",
-    "model = AutoModelForCausalLM.from_pretrained(\n",
-    "    model_name,\n",
-    "    torch_dtype=torch.bfloat16,\n",
-    "    device_map=\"auto\",\n",
-    ")\n",
-    "\n",
-    "lora_config = LoraConfig(\n",
-    "    r=32,\n",
-    "    lora_alpha=32,\n",
-    "    lora_dropout=0,\n",
-    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
-    "                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
-    "    task_type=\"CAUSAL_LM\",\n",
-    ")\n",
-    "\n",
-    "model = get_peft_model(model, lora_config)\n",
-    "model.print_trainable_parameters()\n"
-   ]
+   "outputs": [],
+   "source": "# Cell 2: Load Qwen2.5-3B-Instruct + LoRA\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\nfrom peft import LoraConfig, get_peft_model\n\nmodel_name = \"Qwen/Qwen2.5-3B-Instruct\"\n\ntokenizer = AutoTokenizer.from_pretrained(model_name)\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_name,\n    torch_dtype=torch.bfloat16,\n    device_map=\"auto\",\n)\n\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=16,\n    lora_dropout=0,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n    task_type=\"CAUSAL_LM\",\n)\n\nmodel = get_peft_model(model, lora_config)\nmodel.print_trainable_parameters()"
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -375,68 +156,12 @@
     "id": "D3HDc8_3FG8L",
     "outputId": "b28882c2-1497-4d3b-98b8-4b4cb03cbb0b"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "Connected! Piece: L\n",
-      "Step OK. Reward: -7, Done: False\n",
-      "Environment connection verified!\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 3: Define Tetris client (standalone, no local package needed)\n",
-    "import json\n",
-    "import websockets.sync.client as ws_client\n",
-    "\n",
-    "TETRIS_URL = \"wss://vortexedsquirrel-tetris-env.hf.space/ws\"\n",
-    "\n",
-    "class TetrisClient:\n",
-    "    \"\"\"Lightweight Tetris env client for Colab.\"\"\"\n",
-    "    def __init__(self, url=TETRIS_URL):\n",
-    "        self.url = url\n",
-    "        self.ws = None\n",
-    "\n",
-    "    def connect(self):\n",
-    "        self.ws = ws_client.connect(self.url, open_timeout=30)\n",
-    "        return self\n",
-    "\n",
-    "    def _send_recv(self, msg):\n",
-    "        self.ws.send(json.dumps(msg))\n",
-    "        return json.loads(self.ws.recv(timeout=30))\n",
-    "\n",
-    "    def reset(self, seed=None):\n",
-    "        data = {\"seed\": seed} if seed else {}\n",
-    "        resp = self._send_recv({\"type\": \"reset\", \"data\": data})\n",
-    "        d = resp[\"data\"]\n",
-    "        obs = d.get(\"observation\", d)\n",
-    "        return {\"observation\": obs, \"reward\": d.get(\"reward\", 0), \"done\": d.get(\"done\", False)}\n",
-    "\n",
-    "    def step(self, action_str):\n",
-    "        resp = self._send_recv({\"type\": \"step\", \"data\": {\"action\": action_str, \"metadata\": {}}})\n",
-    "        d = resp[\"data\"]\n",
-    "        obs = d.get(\"observation\", d)\n",
-    "        return {\"observation\": obs, \"reward\": d.get(\"reward\", 0), \"done\": d.get(\"done\", False)}\n",
-    "\n",
-    "    def close(self):\n",
-    "        if self.ws:\n",
-    "            self.ws.close()\n",
-    "\n",
-    "# Quick test\n",
-    "client = TetrisClient().connect()\n",
-    "r = client.reset(seed=42)\n",
-    "print(f\"Connected! Piece: {r['observation']['current_piece']}\")\n",
-    "r = client.step(\"drop\")\n",
-    "print(f\"Step OK. Reward: {r['reward']}, Done: {r['done']}\")\n",
-    "client.close()\n",
-    "print(\"Environment connection verified!\")"
-   ]
+   "outputs": [],
+   "source": "# Cell 3: Game engine + constants + prompt builder\nimport random\nimport torch.nn.functional as F\n\n# Download game engine from repo\n!wget -q -O game_engine.py https://raw.githubusercontent.com/OutOfMystic/tetris-openenv/main/src/tetris_env/server/game_engine.py\nfrom game_engine import TetrisEnv\n\n# === Constants ===\nMAX_ACTIONS_PER_PIECE = 20   # max tokens model can output per piece\nMAX_STEPS_PER_GAME = 200     # max total actions per game\nGAMES_PER_ITER = 8           # games per training iteration (same seed)\nNUM_ITERATIONS = 100         # total training iterations\nTEMPERATURE = 0.7\n\n# Training reward modifiers (on top of engine per-step rewards)\nLR_PENALTY = -0.1            # per L/R move\nPIECE_PLACED_BONUS = 1.0     # per piece successfully placed\nNO_PLACE_PENALTY = -10.0     # if 20 tokens exhausted without placing\nLINE_CLEAR_BONUS = 100.0     # per line cleared (+100/+200/+300/+400)\n\n# Action mapping\nACTION_CHARS = ['L', 'R', 'C', 'W', 'D', 'S']\nACTION_TO_ENGINE = {\n    'L': 'left', 'R': 'right', 'C': 'rotate_cw',\n    'W': 'rotate_ccw', 'D': 'drop', 'S': 'down'\n}\n\n# Pre-compute token IDs for action characters\nACTION_TOKEN_IDS = []\nfor ch in ACTION_CHARS:\n    ids = tokenizer.encode(ch, add_special_tokens=False)\n    ACTION_TOKEN_IDS.append(ids[0])\n    print(f\"  '{ch}' -> token_id {ids[0]}\")\nACTION_TOKEN_IDS = torch.tensor(ACTION_TOKEN_IDS, device=model.device)\n\n# Token ID -> action char lookup\nTOKEN_TO_CHAR = {ACTION_TOKEN_IDS[i].item(): ACTION_CHARS[i] for i in range(6)}\n\nSYSTEM_PROMPT = \"\"\"You are a Tetris AI. You see the board and current piece.\nOutput actions as single letters: L=left R=right C=rotate_cw W=rotate_ccw D=drop S=down\nPlace the piece to fill complete rows. Drop when positioned.\"\"\"\n\ndef build_prompt(result):\n    return f\"\"\"Board:\n{result['board']}\n\nPiece: {result['current_piece']} Next: {result['next_piece']}\nScore: {result['score']} Lines: {result['total_lines']} Height: {result['max_height']} Holes: {result['holes']}\n\nYour actions:\"\"\"\n\nprint(\"\\nGame engine loaded. Action tokens mapped.\")"
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -444,66 +169,12 @@
     "id": "3O3yawZBFG8M",
     "outputId": "7d8ed74b-4c46-4857-c1b7-9bb99848a93f"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "Generated 400 prompts\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 4: Generate training prompts — single-char actions, 100 fixed\n",
-    "import random\n",
-    "\n",
-    "!wget -q -O game_engine.py https://raw.githubusercontent.com/OutOfMystic/tetris-openenv/main/src/tetris_env/server/game_engine.py\n",
-    "\n",
-    "from game_engine import TetrisEnv\n",
-    "\n",
-    "# Action encoding: 1 char = 1 action\n",
-    "ACTION_MAP = {'L': 'left', 'R': 'right', 'C': 'rotate_cw', 'W': 'rotate_ccw', 'D': 'drop', 'S': 'down'}\n",
-    "ACTION_CHARS = list(ACTION_MAP.keys())\n",
-    "\n",
-    "SYSTEM_PROMPT = \"\"\"You are a Tetris AI. Output exactly 100 actions as single letters separated by spaces.\n",
-    "L=left R=right C=rotate_cw W=rotate_ccw D=drop S=down\n",
-    "Example: L L C D R R D L D S D L C D R D ...\n",
-    "Strategy: fill complete rows. Multiple lines at once = bonus.\n",
-    "Output ONLY the 100 letters, nothing else.\"\"\"\n",
-    "\n",
-    "def make_prompt(result):\n",
-    "    return f\"\"\"Board:\n",
-    "{result['board']}\n",
-    "\n",
-    "Piece: {result['current_piece']} Next: {result['next_piece']}\n",
-    "Score: {result['score']} Lines: {result['total_lines']} Height: {result['max_height']} Holes: {result['holes']}\n",
-    "\n",
-    "Your 100 actions:\"\"\"\n",
-    "\n",
-    "def generate_training_prompts(n_prompts=400):\n",
-    "    prompts = []\n",
-    "    for i in range(n_prompts):\n",
-    "        env = TetrisEnv(seed=i)\n",
-    "        result = env.reset(seed=i)\n",
-    "        rng = random.Random(i)\n",
-    "        moves = rng.randint(0, 4)\n",
-    "        direction = rng.choice([\"left\", \"right\"])\n",
-    "        for _ in range(moves):\n",
-    "            env.step(direction)\n",
-    "        result = env._make_result(0)\n",
-    "        prompts.append({\n",
-    "            \"prompt\": [{\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "                       {\"role\": \"user\", \"content\": make_prompt(result)}],\n",
-    "        })\n",
-    "    print(f\"Generated {len(prompts)} prompts\")\n",
-    "    return prompts\n",
-    "\n",
-    "train_prompts = generate_training_prompts(400)\n"
-   ]
+   "outputs": [],
+   "source": "# Cell 4: Core training functions — play_one_game + train_one_iteration\n\ndef play_one_game(model, tokenizer, seed, temperature=TEMPERATURE):\n    \"\"\"\n    Play a full Tetris game. Model sees board before each piece,\n    generates up to 20 action tokens per piece via autoregressive sampling.\n    Returns reward + piece data for gradient computation.\n    \"\"\"\n    env = TetrisEnv(seed=seed)\n    env.reset(seed=seed)\n\n    # Initial random offset (same as old prompt generation)\n    rng = random.Random(seed)\n    moves = rng.randint(0, 4)\n    direction = rng.choice([\"left\", \"right\"])\n    for _ in range(moves):\n        if env.done:\n            break\n        env.step(direction)\n\n    total_reward = 0.0\n    total_steps = 0\n    pieces_data = []\n\n    while not env.done and total_steps < MAX_STEPS_PER_GAME:\n        current_piece_name = env.current_piece_name\n        lines_before = env.total_lines\n\n        # Build fresh prompt for this piece (no history)\n        result = env._make_result(0)\n        messages = [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\", \"content\": build_prompt(result)},\n        ]\n        prompt_ids = tokenizer.apply_chat_template(\n            messages, return_tensors=\"pt\", add_generation_prompt=True\n        ).to(model.device)\n\n        # Autoregressive generation with KV cache\n        input_ids = prompt_ids\n        action_ids_list = []\n        piece_placed = False\n        past_kv = None\n\n        with torch.no_grad():\n            for _ in range(MAX_ACTIONS_PER_PIECE):\n                if past_kv is None:\n                    out = model(input_ids, use_cache=True)\n                    past_kv = out.past_key_values\n                else:\n                    out = model(input_ids[:, -1:], past_key_values=past_kv, use_cache=True)\n                    past_kv = out.past_key_values\n\n                logits = out.logits[:, -1, :]\n\n                # Mask to only 6 action tokens\n                masked = torch.full_like(logits, float('-inf'))\n                masked[0, ACTION_TOKEN_IDS] = logits[0, ACTION_TOKEN_IDS]\n                probs = F.softmax(masked / temperature, dim=-1)\n\n                token_id = torch.multinomial(probs, 1).item()\n                action_ids_list.append(token_id)\n\n                # Execute in engine\n                action_char = TOKEN_TO_CHAR[token_id]\n                action_name = ACTION_TO_ENGINE[action_char]\n                step_result = env.step(action_name)\n\n                total_reward += step_result['reward']\n                total_steps += 1\n\n                if action_name in ('left', 'right'):\n                    total_reward += LR_PENALTY\n\n                # Piece placed? (new piece spawned)\n                if env.current_piece_name != current_piece_name:\n                    piece_placed = True\n                    total_reward += PIECE_PLACED_BONUS\n                    lines_cleared = env.total_lines - lines_before\n                    if lines_cleared > 0:\n                        total_reward += lines_cleared * LINE_CLEAR_BONUS\n                    break\n\n                if env.done:\n                    break\n\n                # Append token for next autoregressive step\n                input_ids = torch.cat([\n                    input_ids,\n                    torch.tensor([[token_id]], device=model.device)\n                ], dim=-1)\n\n        # Force drop if piece not placed in 20 tokens\n        if not piece_placed and not env.done:\n            env.step('drop')\n            total_reward += NO_PLACE_PENALTY\n            total_steps += 1\n            lines_cleared = env.total_lines - lines_before\n            if lines_cleared > 0:\n                total_reward += lines_cleared * LINE_CLEAR_BONUS\n\n        # Store for gradient computation\n        if action_ids_list:\n            pieces_data.append({\n                'prompt_ids': prompt_ids.cpu(),\n                'action_ids': torch.tensor(action_ids_list, dtype=torch.long),\n            })\n\n    return {\n        'reward': total_reward,\n        'pieces': pieces_data,\n        'total_steps': total_steps,\n        'total_lines': env.total_lines,\n        'pieces_placed': len(pieces_data),\n    }\n\n\ndef train_one_iteration(model, optimizer, seed, temperature=TEMPERATURE):\n    \"\"\"\n    One GRPO iteration:\n    1. Rollout: play 8 games (same seed) without grad\n    2. Compute GRPO advantages from game rewards\n    3. Update: recompute log_probs with grad, policy gradient backward\n    \"\"\"\n    # Phase 1: Rollout\n    games = []\n    for _ in range(GAMES_PER_ITER):\n        game = play_one_game(model, tokenizer, seed, temperature)\n        games.append(game)\n\n    rewards = torch.tensor([g['reward'] for g in games], dtype=torch.float32)\n    mean_r = rewards.mean().item()\n    std_r = rewards.std().item()\n\n    # Phase 2: Advantages\n    if std_r < 1e-8:\n        return {'mean_reward': mean_r, 'std_reward': 0.0, 'loss': 0.0,\n                'avg_steps': sum(g['total_steps'] for g in games) / GAMES_PER_ITER,\n                'avg_lines': sum(g['total_lines'] for g in games) / GAMES_PER_ITER,\n                'avg_pieces': sum(g['pieces_placed'] for g in games) / GAMES_PER_ITER}\n\n    advantages = ((rewards - rewards.mean()) / (rewards.std() + 1e-8)).tolist()\n\n    # Phase 3: Update — recompute log_probs with grad, backward per piece\n    optimizer.zero_grad()\n    total_pieces = max(1, sum(len(g['pieces']) for g in games))\n    total_loss = 0.0\n\n    for game_idx, game in enumerate(games):\n        adv = advantages[game_idx]\n        for piece in game['pieces']:\n            prompt = piece['prompt_ids'].to(model.device)\n            actions = piece['action_ids'].to(model.device)\n            if len(actions) == 0:\n                continue\n\n            # Teacher-forced forward pass: prompt + action tokens\n            full_input = torch.cat([prompt.squeeze(0), actions]).unsqueeze(0)\n            logits = model(full_input).logits\n\n            # Log_probs at positions where actions were generated\n            P = prompt.shape[-1]\n            action_logits = logits[0, P-1 : P-1+len(actions), :]\n\n            # Mask to 6 action tokens, apply temperature\n            masked = torch.full_like(action_logits, float('-inf'))\n            masked[:, ACTION_TOKEN_IDS] = action_logits[:, ACTION_TOKEN_IDS]\n            log_probs = F.log_softmax(masked / temperature, dim=-1)\n\n            selected = log_probs.gather(1, actions.unsqueeze(1)).squeeze(1)\n            piece_loss = -(selected.sum() * adv) / total_pieces\n\n            piece_loss.backward()\n            total_loss += piece_loss.item()\n\n    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n    optimizer.step()\n\n    return {\n        'mean_reward': mean_r,\n        'std_reward': std_r,\n        'loss': total_loss,\n        'avg_steps': sum(g['total_steps'] for g in games) / GAMES_PER_ITER,\n        'avg_lines': sum(g['total_lines'] for g in games) / GAMES_PER_ITER,\n        'avg_pieces': sum(g['pieces_placed'] for g in games) / GAMES_PER_ITER,\n    }\n\n# Smoke test\nprint(\"Smoke test: playing 1 game...\")\ntest_game = play_one_game(model, tokenizer, seed=0)\nprint(f\"Reward: {test_game['reward']:.1f}, Steps: {test_game['total_steps']}, \"\n      f\"Pieces: {test_game['pieces_placed']}, Lines: {test_game['total_lines']}\")\nprint(\"Training functions ready.\")"
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -511,61 +182,19 @@
     "id": "JdcKswz5FG8M",
     "outputId": "b0533b78-b10a-4d09-bea1-1b18d8a58c54"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "Dataset size: 400\n",
-      "Example prompt (user message):\n",
-      "Board:\n",
-      "+----------+\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|........@.|\n",
-      "|........@.|\n",
-      "|.......@@.|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "|..........|\n",
-      "+----------+\n",
-      "\n",
-      "Piece:\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 5: Create HF Dataset from prompts\n",
-    "from datasets import Dataset\n",
-    "\n",
-    "dataset = Dataset.from_list(train_prompts)\n",
-    "print(f\"Dataset size: {len(dataset)}\")\n",
-    "print(f\"Example prompt (user message):\")\n",
-    "print(dataset[0]['prompt'][1]['content'][:300])"
-   ]
+   "outputs": [],
+   "source": "# Cell 5: Demo — UNTRAINED model plays one game\nprint(\"=== UNTRAINED MODEL ===\\n\")\n\ngame = play_one_game(model, tokenizer, seed=42)\n\nprint(f\"Total steps: {game['total_steps']}/{MAX_STEPS_PER_GAME}\")\nprint(f\"Pieces placed: {game['pieces_placed']}\")\nprint(f\"Lines cleared: {game['total_lines']}\")\nprint(f\"Game reward: {game['reward']:+.1f}\")\n\n# Show final board by replaying\nenv_demo = TetrisEnv(seed=42)\nenv_demo.reset(seed=42)\nrng = random.Random(42)\nm = rng.randint(0, 4)\nd = rng.choice([\"left\", \"right\"])\nfor _ in range(m):\n    if not env_demo.done:\n        env_demo.step(d)\nfor piece in game['pieces']:\n    for tid in piece['action_ids'].tolist():\n        if not env_demo.done:\n            env_demo.step(ACTION_TO_ENGINE[TOKEN_TO_CHAR[tid]])\n    # Force drop for pieces that weren't placed\n    if not env_demo.done and env_demo.current_piece_name == env_demo.current_piece_name:\n        pass  # engine handles lock internally via gravity\n\nprint(f\"\\nFinal board:\")\nprint(env_demo.board_to_text())\n\nuntrained_reward = game['reward']\nprint(f\"\\nUntrained reward: {untrained_reward:+.1f}\")"
   },
   {
    "cell_type": "code",
-   "source": "# Cell 5b: Demo — UNTRAINED model plays one game (before training)\n# Let's see what the model does BEFORE any GRPO training\n\nimport torch\n\nNUM_ACTIONS = 100  # must match Cell 6 reward function\n\ndef play_one_game(current_model, current_tokenizer, seed=42, label=\"Model\"):\n    \"\"\"Generate 100 actions with the model, play them on local engine, show results.\"\"\"\n    # Create a fresh game\n    env = TetrisEnv(seed=seed)\n    result = env.reset(seed=seed)\n    board_text = make_prompt(result)\n\n    messages = [\n        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n        {\"role\": \"user\", \"content\": board_text},\n    ]\n\n    inputs = current_tokenizer.apply_chat_template(\n        messages, return_tensors=\"pt\", add_generation_prompt=True\n    ).to(current_model.device)\n\n    print(f\"=== {label} playing Tetris (seed={seed}) ===\")\n    print(f\"Starting piece: {result['current_piece']}, Next: {result['next_piece']}\")\n    print(f\"Generating {NUM_ACTIONS} actions...\")\n\n    with torch.no_grad():\n        outputs = current_model.generate(\n            inputs,\n            max_new_tokens=220,\n            temperature=0.7,\n            do_sample=True,\n            pad_token_id=current_tokenizer.pad_token_id,\n        )\n\n    response = current_tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)\n    print(f\"\\nRaw model output:\\n{response[:300]}\")\n\n    # Parse actions\n    chars = [ch for ch in response.strip().upper().split() if ch in ACTION_MAP]\n    if len(chars) < 5:\n        chars = [ch for ch in response.strip().upper() if ch in ACTION_MAP]\n\n    print(f\"\\nParsed {len(chars)} valid actions out of {NUM_ACTIONS} needed\")\n\n    # Pad or truncate\n    if len(chars) < NUM_ACTIONS:\n        chars += ['D'] * (NUM_ACTIONS - len(chars))\n    chars = chars[:NUM_ACTIONS]\n\n    print(f\"Action sequence: {' '.join(chars[:30])} ... {' '.join(chars[-10:])}\")\n\n    # Play on engine\n    env2 = TetrisEnv(seed=seed)\n    env2.reset(seed=seed)\n    steps_played = 0\n    game_over = False\n\n    for ch in chars:\n        r = env2.step(ACTION_MAP[ch])\n        steps_played += 1\n        if r['done']:\n            game_over = True\n            break\n\n    # Calculate training reward (same formula as reward function)\n    total_lines = r['total_lines']\n    holes = r['holes']\n    height = r['max_height']\n    reward = total_lines * 200.0 + steps_played * 0.5 - holes * 2.0 - height * 1.0\n    if game_over:\n        reward -= 50.0\n\n    print(f\"\\n--- Result ---\")\n    print(f\"Steps played: {steps_played}/{NUM_ACTIONS} {'(GAME OVER!)' if game_over else '(survived)'}\")\n    print(f\"Lines cleared: {total_lines}\")\n    print(f\"Score: {r['score']}\")\n    print(f\"Max height: {height}\")\n    print(f\"Holes: {holes}\")\n    print(f\"Training reward: {reward:+.1f}\")\n    print(f\"\\nFinal board:\")\n    print(r['board'])\n\n    return reward\n\n# Run demo with untrained model\nuntrained_reward = play_one_game(model, tokenizer, seed=42, label=\"UNTRAINED model\")\nprint(f\"\\n{'='*50}\")\nprint(f\"Untrained model reward: {untrained_reward:+.1f}\")\nprint(f\"(Remember: higher is better, positive = good play)\")",
+   "source": "# Cell 6: Training loop — 100 iterations of per-piece GRPO\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad],\n    lr=5e-6,\n    weight_decay=0.01,\n)\n\nhistory = []\n\nprint(\"Starting per-piece GRPO training...\")\nprint(f\"Config: {GAMES_PER_ITER} games/iter, max {MAX_STEPS_PER_GAME} steps/game, \"\n      f\"max {MAX_ACTIONS_PER_PIECE} tokens/piece, {NUM_ITERATIONS} iterations\\n\")\n\nfor iteration in range(NUM_ITERATIONS):\n    stats = train_one_iteration(model, optimizer, seed=iteration)\n    history.append(stats)\n\n    if iteration % 5 == 0 or iteration == NUM_ITERATIONS - 1:\n        print(f\"[Iter {iteration:3d}] \"\n              f\"reward={stats['mean_reward']:+8.1f} \"\n              f\"std={stats['std_reward']:6.1f} \"\n              f\"loss={stats['loss']:7.3f} \"\n              f\"steps={stats['avg_steps']:5.1f} \"\n              f\"lines={stats['avg_lines']:4.1f} \"\n              f\"pieces={stats['avg_pieces']:4.1f}\")\n\nprint(\"\\nTraining complete!\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -573,70 +202,8 @@
     "id": "1SRpxwy-FG8M",
     "outputId": "30c168c0-faf0-4c26-fd74-62d3ed3ab13b"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "Test rewards: [-157.5, -138.5, -94.0]\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 6: Reward function — parse single chars, exactly 100 actions\n",
-    "\n",
-    "NUM_ACTIONS = 100\n",
-    "\n",
-    "def tetris_reward_func(completions, **kwargs):\n",
-    "    rewards = []\n",
-    "    for i, completion in enumerate(completions):\n",
-    "        text = completion[0]['content'].strip().upper() if isinstance(completion, list) else str(completion).strip().upper()\n",
-    "\n",
-    "        # Parse single-char actions\n",
-    "        chars = [ch for ch in text.split() if ch in ACTION_MAP]\n",
-    "        # Also try without spaces (e.g. \"LLCDRRД\")\n",
-    "        if len(chars) < 5:\n",
-    "            chars = [ch for ch in text if ch in ACTION_MAP]\n",
-    "\n",
-    "        # Pad to 100 or truncate to 100\n",
-    "        if len(chars) < NUM_ACTIONS:\n",
-    "            chars += ['D'] * (NUM_ACTIONS - len(chars))  # pad with drop\n",
-    "        chars = chars[:NUM_ACTIONS]\n",
-    "\n",
-    "        # Play on engine\n",
-    "        env = TetrisEnv(seed=i)\n",
-    "        env.reset(seed=i)\n",
-    "\n",
-    "        total_lines = 0\n",
-    "        steps_played = 0\n",
-    "        game_over = False\n",
-    "\n",
-    "        for ch in chars:\n",
-    "            result = env.step(ACTION_MAP[ch])\n",
-    "            steps_played += 1\n",
-    "            total_lines = result['total_lines']\n",
-    "            if result['done']:\n",
-    "                game_over = True\n",
-    "                break\n",
-    "\n",
-    "        reward = 0.0\n",
-    "        reward += total_lines * 200.0\n",
-    "        reward += steps_played * 0.5\n",
-    "        reward -= result['holes'] * 2.0\n",
-    "        reward -= result['max_height'] * 1.0\n",
-    "        if game_over:\n",
-    "            reward -= 50.0\n",
-    "\n",
-    "        rewards.append(reward)\n",
-    "\n",
-    "    return rewards\n",
-    "\n",
-    "# Test\n",
-    "test = [[{\"content\": \"L L C D R R D L D S D \" * 6}],\n",
-    "        [{\"content\": \"D D D D D D D D D D \" * 6}],\n",
-    "        [{\"content\": \"garbage text no actions\"}]]\n",
-    "print(f\"Test rewards: {tetris_reward_func(test)}\")\n"
-   ]
+   "outputs": [],
+   "source": "# Cell 7: Plot reward curve\nimport matplotlib.pyplot as plt\n\nfig, axes = plt.subplots(1, 3, figsize=(18, 5))\n\niters = range(len(history))\n\naxes[0].plot(iters, [h['mean_reward'] for h in history])\naxes[0].set_title('Mean Reward per Iteration')\naxes[0].set_xlabel('Iteration')\naxes[0].set_ylabel('Reward')\n\naxes[1].plot(iters, [h['avg_lines'] for h in history])\naxes[1].set_title('Avg Lines Cleared')\naxes[1].set_xlabel('Iteration')\n\naxes[2].plot(iters, [h['loss'] for h in history])\naxes[2].set_title('Policy Gradient Loss')\naxes[2].set_xlabel('Iteration')\n\nplt.tight_layout()\nplt.savefig('reward_curve.png', dpi=150)\nplt.show()\nprint(\"Reward curve saved to reward_curve.png\")"
   },
   {
    "cell_type": "code",
@@ -649,350 +216,15 @@
     "id": "UKvvm6gPFG8M",
     "outputId": "413432d5-f511-40d5-eb43-0874012e593d"
    },
-   "outputs": [
-    {
-     "output_type": "stream",
-     "name": "stderr",
-     "text": [
-      "The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.\n"
-     ]
-    },
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "Starting GRPO training...\n"
-     ]
-    },
-    {
-     "output_type": "stream",
-     "name": "stderr",
-     "text": [
-      "Passing `generation_config` together with generation-related arguments=({'disable_compile'}) is deprecated and will be removed in future versions. Please pass either a `generation_config` object OR all generation parameters explicitly, but not both.\n"
-     ]
-    },
-    {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ],
-      "text/html": [
-       "\n",
-       "    <div>\n",
-       "      \n",
-       "      <progress value='404' max='1200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
-       "      [ 404/1200 1:56:25 < 3:50:31, 0.06 it/s, Epoch 1.01/3]\n",
-       "    </div>\n",
-       "    <table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       " <tr style=\"text-align: left;\">\n",
-       "      <th>Step</th>\n",
-       "      <th>Training Loss</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <td>10</td>\n",
-       "      <td>-0.034254</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>20</td>\n",
-       "      <td>-0.039996</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>30</td>\n",
-       "      <td>-0.005777</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>40</td>\n",
-       "      <td>-0.003115</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>50</td>\n",
-       "      <td>-0.006602</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>60</td>\n",
-       "      <td>-0.034982</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>70</td>\n",
-       "      <td>-0.017585</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>80</td>\n",
-       "      <td>-0.069048</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>90</td>\n",
-       "      <td>-0.050287</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>100</td>\n",
-       "      <td>0.014232</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>110</td>\n",
-       "      <td>0.000296</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>120</td>\n",
-       "      <td>-0.000809</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>130</td>\n",
-       "      <td>0.006139</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>140</td>\n",
-       "      <td>0.002477</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>150</td>\n",
-       "      <td>-0.012867</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>160</td>\n",
-       "      <td>0.003067</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>170</td>\n",
-       "      <td>-0.006496</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>180</td>\n",
-       "      <td>0.006063</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>190</td>\n",
-       "      <td>0.003406</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>200</td>\n",
-       "      <td>-0.022040</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>210</td>\n",
-       "      <td>-0.025758</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>220</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>230</td>\n",
-       "      <td>-0.012708</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>240</td>\n",
-       "      <td>-0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>250</td>\n",
-       "      <td>0.001564</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>260</td>\n",
-       "      <td>-0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>270</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>280</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>290</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>300</td>\n",
-       "      <td>-0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>310</td>\n",
-       "      <td>-0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>320</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>330</td>\n",
-       "      <td>-0.001915</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>340</td>\n",
-       "      <td>-0.001833</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>350</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>360</td>\n",
-       "      <td>0.004288</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>370</td>\n",
-       "      <td>-0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>380</td>\n",
-       "      <td>0.003926</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>390</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <td>400</td>\n",
-       "      <td>0.000000</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table><p>"
-      ]
-     },
-     "metadata": {}
-    },
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "[Step 10] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.725, completions/max_length=220.000, completions/max_terminated_length=150.300, completions/mean_length=192.387, completions/mean_terminated_length=118.600, completions/min_length=88.300, completions/min_terminated_length=88.300, entropy=0.050, frac_reward_zero_std=0.000, grad_norm=0.105, num_tokens=31799.000, reward=-92.294, reward_std=55.018, rewards/tetris_reward_func/mean=-92.294, rewards/tetris_reward_func/std=55.018, step_time=17.325\n",
-      "[Step 20] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.762, completions/max_length=220.000, completions/max_terminated_length=109.700, completions/mean_length=194.350, completions/mean_terminated_length=90.833, completions/min_length=115.200, completions/min_terminated_length=71.200, entropy=0.049, frac_reward_zero_std=0.000, grad_norm=0.154, num_tokens=63787.000, reward=-87.744, reward_std=64.287, rewards/tetris_reward_func/mean=-87.744, rewards/tetris_reward_func/std=64.287, step_time=17.077\n",
-      "[Step 30] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.688, completions/max_length=220.000, completions/max_terminated_length=132.100, completions/mean_length=185.338, completions/mean_terminated_length=113.600, completions/min_length=96.500, completions/min_terminated_length=96.500, entropy=0.051, frac_reward_zero_std=0.000, grad_norm=0.127, num_tokens=94966.000, reward=-84.419, reward_std=60.849, rewards/tetris_reward_func/mean=-84.419, rewards/tetris_reward_func/std=60.849, step_time=17.085\n",
-      "[Step 40] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.725, completions/max_length=220.000, completions/max_terminated_length=114.200, completions/mean_length=189.113, completions/mean_terminated_length=100.702, completions/min_length=107.200, completions/min_terminated_length=85.200, entropy=0.049, frac_reward_zero_std=0.000, grad_norm=0.059, num_tokens=126439.000, reward=-82.194, reward_std=65.590, rewards/tetris_reward_func/mean=-82.194, rewards/tetris_reward_func/std=65.590, step_time=17.078\n",
-      "[Step 50] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.738, completions/max_length=220.000, completions/max_terminated_length=128.500, completions/mean_length=193.863, completions/mean_terminated_length=111.117, completions/min_length=120.200, completions/min_terminated_length=98.200, entropy=0.041, frac_reward_zero_std=0.000, grad_norm=0.061, num_tokens=158364.000, reward=-72.206, reward_std=68.607, rewards/tetris_reward_func/mean=-72.206, rewards/tetris_reward_func/std=68.607, step_time=17.067\n",
-      "[Step 60] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.800, completions/max_length=220.000, completions/max_terminated_length=133.100, completions/mean_length=200.500, completions/mean_terminated_length=109.633, completions/min_length=111.000, completions/min_terminated_length=89.000, entropy=0.038, frac_reward_zero_std=0.000, grad_norm=0.059, num_tokens=190708.000, reward=-66.744, reward_std=67.078, rewards/tetris_reward_func/mean=-66.744, rewards/tetris_reward_func/std=67.078, step_time=17.116\n",
-      "[Step 70] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.800, completions/max_length=220.000, completions/max_terminated_length=121.900, completions/mean_length=200.662, completions/mean_terminated_length=111.100, completions/min_length=119.600, completions/min_terminated_length=97.600, entropy=0.039, frac_reward_zero_std=0.000, grad_norm=0.107, num_tokens=223145.000, reward=-46.081, reward_std=68.570, rewards/tetris_reward_func/mean=-46.081, rewards/tetris_reward_func/std=68.570, step_time=17.263\n",
-      "[Step 80] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.738, completions/max_length=220.000, completions/max_terminated_length=127.900, completions/mean_length=193.125, completions/mean_terminated_length=111.063, completions/min_length=96.800, completions/min_terminated_length=96.800, entropy=0.052, frac_reward_zero_std=0.000, grad_norm=0.153, num_tokens=254939.000, reward=-19.050, reward_std=55.485, rewards/tetris_reward_func/mean=-19.050, rewards/tetris_reward_func/std=55.485, step_time=17.245\n",
-      "[Step 90] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.750, completions/max_length=220.000, completions/max_terminated_length=119.900, completions/mean_length=197.775, completions/mean_terminated_length=103.608, completions/min_length=128.900, completions/min_terminated_length=84.900, entropy=0.052, frac_reward_zero_std=0.000, grad_norm=0.117, num_tokens=287121.000, reward=6.112, reward_std=42.999, rewards/tetris_reward_func/mean=6.112, rewards/tetris_reward_func/std=42.999, step_time=17.200\n",
-      "[Step 100] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.900, completions/max_length=220.000, completions/max_terminated_length=61.200, completions/mean_length=209.688, completions/mean_terminated_length=59.300, completions/min_length=167.400, completions/min_terminated_length=57.400, entropy=0.042, frac_reward_zero_std=0.000, grad_norm=0.099, num_tokens=320240.000, reward=6.544, reward_std=43.744, rewards/tetris_reward_func/mean=6.544, rewards/tetris_reward_func/std=43.744, step_time=17.200\n",
-      "[Step 110] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.988, completions/max_length=220.000, completions/max_terminated_length=16.400, completions/mean_length=219.300, completions/mean_terminated_length=16.400, completions/min_length=214.400, completions/min_terminated_length=16.400, entropy=0.035, frac_reward_zero_std=0.000, grad_norm=0.117, num_tokens=354208.000, reward=14.625, reward_std=26.601, rewards/tetris_reward_func/mean=14.625, rewards/tetris_reward_func/std=26.601, step_time=17.167\n",
-      "[Step 120] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.938, completions/max_length=220.000, completions/max_terminated_length=60.600, completions/mean_length=215.675, completions/mean_terminated_length=57.950, completions/min_length=187.300, completions/min_terminated_length=55.300, entropy=0.035, frac_reward_zero_std=0.000, grad_norm=0.119, num_tokens=387822.000, reward=16.738, reward_std=15.117, rewards/tetris_reward_func/mean=16.738, rewards/tetris_reward_func/std=15.117, step_time=17.196\n",
-      "[Step 130] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.950, completions/max_length=220.000, completions/max_terminated_length=61.300, completions/mean_length=216.662, completions/mean_terminated_length=61.300, completions/min_length=193.300, completions/min_terminated_length=61.300, entropy=0.035, frac_reward_zero_std=0.000, grad_norm=0.103, num_tokens=421579.000, reward=16.794, reward_std=25.586, rewards/tetris_reward_func/mean=16.794, rewards/tetris_reward_func/std=25.586, step_time=17.231\n",
-      "[Step 140] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.925, completions/max_length=220.000, completions/max_terminated_length=82.500, completions/mean_length=215.450, completions/mean_terminated_length=80.400, completions/min_length=188.300, completions/min_terminated_length=78.300, entropy=0.041, frac_reward_zero_std=0.000, grad_norm=0.142, num_tokens=455215.000, reward=24.444, reward_std=27.435, rewards/tetris_reward_func/mean=24.444, rewards/tetris_reward_func/std=27.435, step_time=17.233\n"
-     ]
-    },
-    {
-     "output_type": "stream",
-     "name": "stderr",
-     "text": [
-      "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n",
-      "WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
-     ]
-    },
-    {
-     "output_type": "stream",
-     "name": "stdout",
-     "text": [
-      "[Step 150] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.963, completions/max_length=220.000, completions/max_terminated_length=23.100, completions/mean_length=215.700, completions/mean_terminated_length=20.800, completions/min_length=194.500, completions/min_terminated_length=18.500, entropy=0.053, frac_reward_zero_std=0.000, grad_norm=0.098, num_tokens=488775.000, reward=26.969, reward_std=42.222, rewards/tetris_reward_func/mean=26.969, rewards/tetris_reward_func/std=42.222, step_time=17.195\n",
-      "[Step 160] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.938, completions/max_length=220.000, completions/max_terminated_length=70.000, completions/mean_length=215.000, completions/mean_terminated_length=70.000, completions/min_length=180.000, completions/min_terminated_length=70.000, entropy=0.061, frac_reward_zero_std=0.000, grad_norm=0.141, num_tokens=522343.000, reward=23.725, reward_std=18.791, rewards/tetris_reward_func/mean=23.725, rewards/tetris_reward_func/std=18.791, step_time=17.401\n",
-      "[Step 170] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.988, completions/max_length=220.000, completions/max_terminated_length=13.200, completions/mean_length=218.900, completions/mean_terminated_length=13.200, completions/min_length=211.200, completions/min_terminated_length=13.200, entropy=0.083, frac_reward_zero_std=0.000, grad_norm=0.217, num_tokens=556311.000, reward=24.812, reward_std=26.831, rewards/tetris_reward_func/mean=24.812, rewards/tetris_reward_func/std=26.831, step_time=17.467\n",
-      "[Step 180] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.988, completions/max_length=220.000, completions/max_terminated_length=14.300, completions/mean_length=219.037, completions/mean_terminated_length=14.300, completions/min_length=212.300, completions/min_terminated_length=14.300, entropy=0.093, frac_reward_zero_std=0.000, grad_norm=0.196, num_tokens=590178.000, reward=27.462, reward_std=19.135, rewards/tetris_reward_func/mean=27.462, rewards/tetris_reward_func/std=19.135, step_time=17.405\n",
-      "[Step 190] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.975, completions/max_length=220.000, completions/max_terminated_length=29.900, completions/mean_length=218.238, completions/mean_terminated_length=29.900, completions/min_length=205.900, completions/min_terminated_length=29.900, entropy=0.119, frac_reward_zero_std=0.000, grad_norm=0.220, num_tokens=623981.000, reward=13.000, reward_std=26.462, rewards/tetris_reward_func/mean=13.000, rewards/tetris_reward_func/std=26.462, step_time=17.372\n",
-      "[Step 200] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.938, completions/max_length=220.000, completions/max_terminated_length=49.800, completions/mean_length=213.275, completions/mean_terminated_length=47.650, completions/min_length=177.500, completions/min_terminated_length=45.500, entropy=0.186, frac_reward_zero_std=0.000, grad_norm=0.273, num_tokens=657443.000, reward=13.613, reward_std=22.587, rewards/tetris_reward_func/mean=13.613, rewards/tetris_reward_func/std=22.587, step_time=17.294\n",
-      "[Step 210] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.863, completions/max_length=220.000, completions/max_terminated_length=98.300, completions/mean_length=208.775, completions/mean_terminated_length=92.817, completions/min_length=154.200, completions/min_terminated_length=88.200, entropy=0.255, frac_reward_zero_std=0.000, grad_norm=0.229, num_tokens=690505.000, reward=18.550, reward_std=19.927, rewards/tetris_reward_func/mean=18.550, rewards/tetris_reward_func/std=19.927, step_time=17.190\n",
-      "[Step 220] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.281, frac_reward_zero_std=0.000, grad_norm=0.288, num_tokens=724465.000, reward=19.150, reward_std=20.560, rewards/tetris_reward_func/mean=19.150, rewards/tetris_reward_func/std=20.560, step_time=17.162\n",
-      "[Step 230] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.975, completions/max_length=220.000, completions/max_terminated_length=25.300, completions/mean_length=217.662, completions/mean_terminated_length=25.300, completions/min_length=201.300, completions/min_terminated_length=25.300, entropy=0.201, frac_reward_zero_std=0.000, grad_norm=0.329, num_tokens=758206.000, reward=18.900, reward_std=18.913, rewards/tetris_reward_func/mean=18.900, rewards/tetris_reward_func/std=18.913, step_time=17.190\n",
-      "[Step 240] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.198, frac_reward_zero_std=0.000, grad_norm=0.214, num_tokens=792254.000, reward=19.625, reward_std=16.034, rewards/tetris_reward_func/mean=19.625, rewards/tetris_reward_func/std=16.034, step_time=17.204\n",
-      "[Step 250] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.975, completions/max_length=220.000, completions/max_terminated_length=35.000, completions/mean_length=218.875, completions/mean_terminated_length=35.000, completions/min_length=211.000, completions/min_terminated_length=35.000, entropy=0.192, frac_reward_zero_std=0.000, grad_norm=0.198, num_tokens=826172.000, reward=25.400, reward_std=11.692, rewards/tetris_reward_func/mean=25.400, rewards/tetris_reward_func/std=11.692, step_time=17.102\n",
-      "[Step 260] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.173, frac_reward_zero_std=0.000, grad_norm=0.223, num_tokens=860124.000, reward=25.200, reward_std=15.597, rewards/tetris_reward_func/mean=25.200, rewards/tetris_reward_func/std=15.597, step_time=17.231\n",
-      "[Step 270] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.221, frac_reward_zero_std=0.000, grad_norm=0.288, num_tokens=894140.000, reward=25.312, reward_std=12.701, rewards/tetris_reward_func/mean=25.312, rewards/tetris_reward_func/std=12.701, step_time=17.220\n",
-      "[Step 280] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.243, frac_reward_zero_std=0.000, grad_norm=0.274, num_tokens=928092.000, reward=26.238, reward_std=12.358, rewards/tetris_reward_func/mean=26.238, rewards/tetris_reward_func/std=12.358, step_time=17.183\n",
-      "[Step 290] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.151, frac_reward_zero_std=0.000, grad_norm=0.191, num_tokens=962060.000, reward=26.725, reward_std=13.212, rewards/tetris_reward_func/mean=26.725, rewards/tetris_reward_func/std=13.212, step_time=17.353\n",
-      "[Step 300] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.135, frac_reward_zero_std=0.000, grad_norm=0.218, num_tokens=996084.000, reward=27.775, reward_std=11.482, rewards/tetris_reward_func/mean=27.775, rewards/tetris_reward_func/std=11.482, step_time=17.287\n",
-      "[Step 310] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.134, frac_reward_zero_std=0.000, grad_norm=0.223, num_tokens=1030036.000, reward=25.488, reward_std=12.503, rewards/tetris_reward_func/mean=25.488, rewards/tetris_reward_func/std=12.503, step_time=17.345\n",
-      "[Step 320] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.148, frac_reward_zero_std=0.000, grad_norm=0.251, num_tokens=1064028.000, reward=23.150, reward_std=15.532, rewards/tetris_reward_func/mean=23.150, rewards/tetris_reward_func/std=15.532, step_time=17.364\n",
-      "[Step 330] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.988, completions/max_length=220.000, completions/max_terminated_length=18.400, completions/mean_length=219.550, completions/mean_terminated_length=18.400, completions/min_length=216.400, completions/min_terminated_length=18.400, entropy=0.179, frac_reward_zero_std=0.000, grad_norm=0.244, num_tokens=1097984.000, reward=24.875, reward_std=14.287, rewards/tetris_reward_func/mean=24.875, rewards/tetris_reward_func/std=14.287, step_time=17.403\n",
-      "[Step 340] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.963, completions/max_length=220.000, completions/max_terminated_length=40.900, completions/mean_length=219.150, completions/mean_terminated_length=40.350, completions/min_length=215.800, completions/min_terminated_length=39.800, entropy=0.239, frac_reward_zero_std=0.000, grad_norm=0.319, num_tokens=1131876.000, reward=26.175, reward_std=12.656, rewards/tetris_reward_func/mean=26.175, rewards/tetris_reward_func/std=12.656, step_time=17.309\n",
-      "[Step 350] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.257, frac_reward_zero_std=0.000, grad_norm=0.404, num_tokens=1165884.000, reward=24.625, reward_std=12.939, rewards/tetris_reward_func/mean=24.625, rewards/tetris_reward_func/std=12.939, step_time=17.213\n",
-      "[Step 360] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.975, completions/max_length=220.000, completions/max_terminated_length=31.200, completions/mean_length=218.400, completions/mean_terminated_length=31.200, completions/min_length=207.200, completions/min_terminated_length=31.200, entropy=0.283, frac_reward_zero_std=0.000, grad_norm=0.297, num_tokens=1199772.000, reward=26.212, reward_std=13.559, rewards/tetris_reward_func/mean=26.212, rewards/tetris_reward_func/std=13.559, step_time=17.201\n",
-      "[Step 370] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.232, frac_reward_zero_std=0.000, grad_norm=0.318, num_tokens=1233708.000, reward=28.762, reward_std=9.307, rewards/tetris_reward_func/mean=28.762, rewards/tetris_reward_func/std=9.307, step_time=17.198\n",
-      "[Step 380] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=0.988, completions/max_length=220.000, completions/max_terminated_length=16.400, completions/mean_length=219.300, completions/mean_terminated_length=16.400, completions/min_length=214.400, completions/min_terminated_length=16.400, entropy=0.198, frac_reward_zero_std=0.000, grad_norm=0.325, num_tokens=1267652.000, reward=31.312, reward_std=9.031, rewards/tetris_reward_func/mean=31.312, rewards/tetris_reward_func/std=9.031, step_time=17.213\n",
-      "[Step 390] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.141, frac_reward_zero_std=0.000, grad_norm=0.231, num_tokens=1301596.000, reward=28.175, reward_std=10.578, rewards/tetris_reward_func/mean=28.175, rewards/tetris_reward_func/std=10.578, step_time=17.176\n",
-      "[Step 400] clip_ratio/high_max=0.000, clip_ratio/high_mean=0.000, clip_ratio/low_mean=0.000, clip_ratio/low_min=0.000, clip_ratio/region_mean=0.000, completions/clipped_ratio=1.000, completions/max_length=220.000, completions/max_terminated_length=0.000, completions/mean_length=220.000, completions/mean_terminated_length=0.000, completions/min_length=220.000, completions/min_terminated_length=0.000, entropy=0.122, frac_reward_zero_std=0.000, grad_norm=0.358, num_tokens=1335572.000, reward=28.887, reward_std=10.549, rewards/tetris_reward_func/mean=28.887, rewards/tetris_reward_func/std=10.549, step_time=17.246\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Cell 7: Configure and run GRPO training (A100, with reward logging)\n",
-    "from trl import GRPOConfig, GRPOTrainer\n",
-    "from transformers import TrainerCallback\n",
-    "\n",
-    "class RewardLogger(TrainerCallback):\n",
-    "    def on_log(self, args, state, control, logs=None, **kwargs):\n",
-    "        if not logs:\n",
-    "            return\n",
-    "        step = state.global_step\n",
-    "        parts = []\n",
-    "        for k in sorted(logs):\n",
-    "            if k != 'loss' and k != 'learning_rate' and k != 'epoch':\n",
-    "                parts.append(f\"{k}={logs[k]:.3f}\")\n",
-    "        if parts:\n",
-    "            print(f\"[Step {step}] {', '.join(parts)}\")\n",
-    "\n",
-    "training_args = GRPOConfig(\n",
-    "    output_dir=\"tetris-agent-grpo\",\n",
-    "    num_train_epochs=3,\n",
-    "    per_device_train_batch_size=4,\n",
-    "    gradient_accumulation_steps=2,\n",
-    "    num_generations=8,\n",
-    "    max_completion_length=220,\n",
-    "    learning_rate=5e-6,\n",
-    "    logging_steps=10,             # log every 10 steps (less spam)\n",
-    "    save_steps=50,\n",
-    "    bf16=True,\n",
-    "    report_to=\"none\",\n",
-    ")\n",
-    "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model=model,\n",
-    "    processing_class=tokenizer,\n",
-    "    reward_funcs=[tetris_reward_func],\n",
-    "    args=training_args,\n",
-    "    train_dataset=dataset,\n",
-    "    callbacks=[RewardLogger()],\n",
-    ")\n",
-    "\n",
-    "print(\"Starting GRPO training...\")\n",
-    "trainer.train()\n",
-    "print(\"Training complete!\")\n"
-   ]
+   "outputs": [],
+   "source": "# Cell 8: Demo — TRAINED model plays 3 games\nprint(\"=== TRAINED MODEL ===\\n\")\n\ntrained_rewards = []\nfor seed in [42, 123, 7]:\n    game = play_one_game(model, tokenizer, seed=seed)\n    print(f\"Seed {seed}: reward={game['reward']:+.1f}, \"\n          f\"steps={game['total_steps']}, lines={game['total_lines']}, \"\n          f\"pieces={game['pieces_placed']}\")\n    trained_rewards.append(game['reward'])\n\navg_trained = sum(trained_rewards) / len(trained_rewards)\nprint(f\"\\n{'='*50}\")\nprint(f\"UNTRAINED reward (seed=42): {untrained_reward:+.1f}\")\nprint(f\"TRAINED avg reward (3 games): {avg_trained:+.1f}\")\nprint(f\"Improvement: {avg_trained - untrained_reward:+.1f}\")\nprint('='*50)\n\nif avg_trained > untrained_reward:\n    print(\"Training improved the model!\")\nelse:\n    print(\"Model needs more training or tuning.\")"
   },
   {
    "cell_type": "code",
-   "source": "# Cell 8: Plot reward curve\nimport matplotlib.pyplot as plt\n\nlogs = trainer.state.log_history\ntrain_logs = [l for l in logs if 'loss' in l]\nreward_logs = [l for l in logs if 'reward' in l or 'rewards/mean' in l]\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\nif train_logs:\n    axes[0].plot([l.get('step', i) for i, l in enumerate(train_logs)],\n                 [l['loss'] for l in train_logs])\n    axes[0].set_title('Training Loss')\n    axes[0].set_xlabel('Step')\n    axes[0].set_ylabel('Loss')\n\nreward_key = 'reward' if reward_logs and 'reward' in reward_logs[0] else 'rewards/mean'\nif reward_logs:\n    axes[1].plot([l.get('step', i) for i, l in enumerate(reward_logs)],\n                 [l.get(reward_key, 0) for l in reward_logs])\n    axes[1].set_title('Mean Reward (should go up!)')\n    axes[1].set_xlabel('Step')\n    axes[1].set_ylabel('Reward')\n\nplt.tight_layout()\nplt.savefig('reward_curve.png', dpi=150)\nplt.show()\nprint(\"Reward curve saved to reward_curve.png\")",
+   "source": "# Cell 9: Push trained model to HF Hub\nmodel.push_to_hub(\"VortexedSquirrel/tetris-agent-grpo\")\ntokenizer.push_to_hub(\"VortexedSquirrel/tetris-agent-grpo\")\nprint(\"Model pushed to https://huggingface.co/VortexedSquirrel/tetris-agent-grpo\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "HUXSqr3nFG8N"
-   },
-   "outputs": [],
-   "source": "# Cell 9: Demo — TRAINED model plays Tetris (after training)\n# Compare with untrained model from Cell 5b!\n\nprint(\"Playing 3 games with TRAINED model and comparing to untrained baseline...\\n\")\n\ntrained_rewards = []\nfor seed in [42, 123, 7]:\n    r = play_one_game(model, tokenizer, seed=seed, label=f\"TRAINED model\")\n    trained_rewards.append(r)\n    print()\n\navg_trained = sum(trained_rewards) / len(trained_rewards)\n\nprint(\"=\" * 60)\nprint(f\"UNTRAINED model reward (seed=42):  {untrained_reward:+.1f}\")\nprint(f\"TRAINED model avg reward (3 games): {avg_trained:+.1f}\")\nprint(f\"Improvement: {avg_trained - untrained_reward:+.1f}\")\nprint(\"=\" * 60)\n\nif avg_trained > untrained_reward:\n    print(\"Training improved the model! GRPO is working.\")\nelse:\n    print(\"Model needs more training epochs or tuning.\")\n"
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "BnOW5pcKFG8N"
-   },
-   "outputs": [],
-   "source": [
-    "# Cell 10: Push trained model to HF Hub\n",
-    "model.push_to_hub(\"VortexedSquirrel/tetris-agent-grpo\")\n",
-    "tokenizer.push_to_hub(\"VortexedSquirrel/tetris-agent-grpo\")\n",
-    "print(\"Model pushed to https://huggingface.co/VortexedSquirrel/tetris-agent-grpo\")"
-   ]
   }
  ],
  "metadata": {