OpenEnv/docs/source/tutorials/rl-training-2048.md

# RL Training with OpenEnv: 2048 Game

This tutorial covers training a language model to play the 2048 game using
reinforcement learning with GRPO (Group Relative Policy Optimization).

```{note}
**Time**: ~45 minutes | **Difficulty**: Advanced | **GPU Required**: Yes (T4 or better)
```

## What You'll Learn

- **Model Setup**: Load and configure LLMs with Unsloth for efficient RL
- **Environment Connection**: Connect to the 2048 OpenEnv environment
- **Reward Design**: Create effective reward functions
- **GRPO Training**: Train models with reinforcement learning
- **Deployment**: Save and deploy trained models

## Prerequisites

Before starting this tutorial, you should have completed the
[Getting Started](/auto_getting_started/index) series to understand:

- How OpenEnv environments work
- The reset/step/state API pattern
- How to connect to environments

You'll also need:

- A GPU (free T4 on Google Colab works)
- Basic understanding of PyTorch
- ~30 minutes for training

## Part 1: Environment Setup

### Installation

```bash
# Install required packages
!pip install -q unsloth openenv-core trl

# For Google Colab, also run:
!pip install -q "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
```

### Imports

```python
import torch
from dataclasses import dataclass
from typing import List, Optional, Dict, Any
import random

# Check GPU availability
print(f"GPU Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
```

## Part 2: Model Configuration

We use Unsloth for memory-efficient training with LoRA adapters.

### Configuration Classes

```python
@dataclass
class ModelConfig:
    """Configuration for loading LLM models."""
    model_name: str = "unsloth/Qwen2.5-1.5B"
    max_seq_length: int = 768
    load_in_4bit: bool = True
    dtype: Optional[str] = None  # Auto-detect


@dataclass
class LoRAConfig:
    """Configuration for LoRA fine-tuning."""
    r: int = 16
    lora_alpha: int = 32
    target_modules: List[str] = None
    lora_dropout: float = 0.0

    def __post_init__(self):
        if self.target_modules is None:
            self.target_modules = [
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj",
            ]
```

### Loading the Model

```python
from unsloth import FastLanguageModel

# Create configurations
model_config = ModelConfig()
lora_config = LoRAConfig()

# Load model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_config.model_name,
    max_seq_length=model_config.max_seq_length,
    load_in_4bit=model_config.load_in_4bit,
    dtype=model_config.dtype,
)

# Apply LoRA adapters
model = FastLanguageModel.get_peft_model(
    model,
    r=lora_config.r,
    target_modules=lora_config.target_modules,
    lora_alpha=lora_config.lora_alpha,
    lora_dropout=lora_config.lora_dropout,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42,
)

# Check parameter counts
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable:,} / {total:,} ({trainable/total*100:.2f}%)")
```

## Part 3: The 2048 Environment

### Game Overview

2048 is a sliding puzzle game where you combine tiles to reach 2048.

**Actions:**
- `0` = UP
- `1` = RIGHT
- `2` = DOWN
- `3` = LEFT

**Goal:** Create a tile with value 2048 (or higher!)

### Connecting to the Environment

```python
from envs.openspiel_env import OpenSpielEnv, OpenSpielAction

# Connect to 2048 environment
# Option 1: From Hub
env = OpenSpielEnv.from_hub("openenv/openspiel-env")

# Option 2: From running server
# env = OpenSpielEnv(base_url="http://localhost:8000")

# Test connection
with env:
    result = env.reset()
    print(f"Game started!")
    print(f"Legal actions: {result.observation.legal_actions}")

    # Take a test action
    action = OpenSpielAction(action_id=0, game_name="2048")
    result = env.step(action)
    print(f"After UP: reward={result.reward}, done={result.done}")
```

### Board Utilities

```python
import numpy as np
from typing import List

def info_state_to_board(info_state: List[int], size: int = 4) -> List[List[int]]:
    """Convert flat info_state to 2D board."""
    return np.array(info_state, dtype=int).reshape(size, size).tolist()

def render_board(board: List[List[int]]) -> str:
    """Render board as ASCII string."""
    lines = ["+------" * len(board[0]) + "+"]
    for row in board:
        cells = [f"{v:5d}" if v > 0 else "    ." for v in row]
        lines.append("|" + " |".join(cells) + " |")
        lines.append("+------" * len(row) + "+")
    return "\n".join(lines)

def get_max_tile(board: List[List[int]]) -> int:
    """Get highest tile value."""
    return max(cell for row in board for cell in row)
```

## Part 4: Reward Function Design

The reward function is crucial for RL. We consider:

1. **Success**: Did we reach 2048?
2. **Progress**: What's the highest tile achieved?
3. **Code Quality**: Did the generated code execute correctly?

### Reward Implementation

```python
import math

def calculate_reward(
    max_tile: int,
    success: bool,
    code_error: bool = False
) -> float:
    """
    Calculate reward for a 2048 game outcome.

    Args:
        max_tile: Highest tile achieved (2, 4, 8, ..., 2048)
        success: Whether we reached 2048
        code_error: Whether generated code had errors

    Returns:
        Float reward value
    """
    if code_error:
        return -0.5  # Penalty for invalid code

    if success:
        return 1.0  # Full reward for winning

    # Progress reward: log scale from 0 to 0.9
    if max_tile > 0:
        progress = math.log2(max_tile) / math.log2(2048)
        return min(0.9, progress)

    return 0.0

# Test reward function
test_cases = [
    (2048, True, False, "Won!"),
    (1024, False, False, "Got to 1024"),
    (512, False, False, "Got to 512"),
    (64, False, False, "Early game"),
]

for max_tile, success, error, desc in test_cases:
    reward = calculate_reward(max_tile, success, error)
    print(f"{desc:20s} -> Reward: {reward:+.3f}")
```

## Part 5: Strategy Generation

We'll train the model to generate Python strategy functions.

### Prompt Template

```python
SYSTEM_PROMPT = """You are an expert at playing 2048. Generate a Python function
that takes a board state and returns the best action (0=UP, 1=RIGHT, 2=DOWN, 3=LEFT).

The board is a 4x4 list of integers. Empty cells are 0.
Your function should analyze the board and return an optimal move.
"""

def create_prompt(board: List[List[int]]) -> str:
    """Create prompt for strategy generation."""
    board_str = "\n".join(str(row) for row in board)
    return f"""{SYSTEM_PROMPT}

Current board:
{board_str}

Generate a strategy function:
```python
def strategy(board):
    # Your code here
    return action  # 0, 1, 2, or 3
```"""
```

### Executing Generated Strategies

```python
import ast
from typing import Callable

def extract_and_execute_strategy(
    generated_code: str,
    board: List[List[int]],
    timeout: float = 5.0
) -> tuple[int, bool]:
    """
    Extract and execute a generated strategy function.

    Returns:
        (action, success): The action to take and whether execution succeeded
    """
    try:
        # Extract code block
        if "```python" in generated_code:
            code = generated_code.split("```python")[1].split("```")[0]
        else:
            code = generated_code

        # Parse and validate AST
        tree = ast.parse(code)

        # Execute in sandbox
        namespace = {"board": board}
        exec(compile(tree, "<strategy>", "exec"), namespace)

        # Call the strategy function
        if "strategy" in namespace:
            action = namespace["strategy"](board)
            if action in [0, 1, 2, 3]:
                return action, True

        return 0, False  # Default action on failure

    except Exception as e:
        print(f"Strategy execution error: {e}")
        return 0, False
```

## Part 6: GRPO Training

GRPO (Group Relative Policy Optimization) is optimized for language models.

### Training Configuration

```python
from trl import GRPOConfig, GRPOTrainer

grpo_config = GRPOConfig(
    # Learning rate
    learning_rate=2e-6,

    # Batch sizes
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,

    # Training duration
    max_steps=200,

    # Memory optimization
    bf16=True,
    gradient_checkpointing=True,

    # Logging
    logging_steps=1,
    output_dir="./2048_grpo_output",
    report_to="none",
)
```

### Training Loop

```python
def train_2048_agent(
    model,
    tokenizer,
    env,
    config: GRPOConfig,
    num_episodes: int = 100,
):
    """
    Train the model to play 2048 using GRPO.
    """
    # Prepare model for training
    FastLanguageModel.for_training(model)

    training_data = []

    for episode in range(num_episodes):
        # Reset environment
        result = env.reset()
        board = info_state_to_board(result.observation.info_state)

        episode_reward = 0
        steps = 0

        while not result.done and steps < 1000:
            # Generate strategy
            prompt = create_prompt(board)
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                temperature=0.7,
                do_sample=True,
            )

            generated = tokenizer.decode(outputs[0], skip_special_tokens=True)

            # Execute strategy
            action, success = extract_and_execute_strategy(generated, board)

            # Take action in environment
            env_action = OpenSpielAction(action_id=action, game_name="2048")
            result = env.step(env_action)

            # Update board
            board = info_state_to_board(result.observation.info_state)
            episode_reward += result.reward if result.reward else 0
            steps += 1

        # Calculate final reward
        max_tile = get_max_tile(board)
        final_reward = calculate_reward(max_tile, max_tile >= 2048)

        # Store for training
        training_data.append({
            "prompt": prompt,
            "response": generated,
            "reward": final_reward,
        })

        if episode % 10 == 0:
            print(f"Episode {episode}: Max tile={max_tile}, Reward={final_reward:.3f}")

    return training_data
```

## Part 7: Deployment

After training, save and deploy your model.

### Saving the Model

```python
# Save LoRA adapters only
model.save_pretrained("./2048_strategy_model")
tokenizer.save_pretrained("./2048_strategy_model")

# Save merged model for inference
model.save_pretrained_merged(
    "./2048_strategy_model_merged",
    tokenizer,
    save_method="merged_16bit",
)
```

### Push to Hugging Face Hub

```python
# Push to Hub
model.push_to_hub(
    "your-username/2048-strategy-model",
    tokenizer,
    save_method="merged_16bit",
    private=False,
)

print("Model deployed to: huggingface.co/your-username/2048-strategy-model")
```

### Using the Trained Model

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load trained model
model = AutoModelForCausalLM.from_pretrained("your-username/2048-strategy-model")
tokenizer = AutoTokenizer.from_pretrained("your-username/2048-strategy-model")

# Generate strategy
def get_action(board: List[List[int]]) -> int:
    prompt = create_prompt(board)
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=256)
    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    action, _ = extract_and_execute_strategy(generated, board)
    return action

# Play a game
with OpenSpielEnv.from_hub("openenv/openspiel-env") as env:
    result = env.reset()
    board = info_state_to_board(result.observation.info_state)

    while not result.done:
        action = get_action(board)
        result = env.step(OpenSpielAction(action_id=action, game_name="2048"))
        board = info_state_to_board(result.observation.info_state)

    print(f"Final max tile: {get_max_tile(board)}")
```

## Preventing Reward Hacking

Be aware of potential reward hacking strategies:

1. **Code that modifies rewards** - Run in sandboxed environment
2. **Infinite loops** - Set execution timeouts
3. **Memory exhaustion** - Limit resource usage

```python
import resource
import signal

def safe_execute(code: str, board: List[List[int]], timeout: float = 5.0) -> int:
    """Execute strategy with safety limits."""

    def handler(signum, frame):
        raise TimeoutError("Strategy timed out")

    # Set timeout
    signal.signal(signal.SIGALRM, handler)
    signal.alarm(int(timeout))

    try:
        # Set memory limit (100MB)
        resource.setrlimit(resource.RLIMIT_AS, (100 * 1024 * 1024, -1))

        # Execute in restricted namespace
        namespace = {"board": board, "__builtins__": {"len": len, "max": max, "min": min}}
        exec(code, namespace)

        return namespace.get("strategy", lambda b: 0)(board)
    finally:
        signal.alarm(0)
```

## Summary

In this tutorial, you learned:

1. **Model Setup**: Loading LLMs with Unsloth and LoRA
2. **Environment Connection**: Using OpenEnv's 2048 environment
3. **Reward Design**: Creating balanced reward functions
4. **GRPO Training**: Training with reinforcement learning
5. **Deployment**: Saving and sharing trained models

## Next Steps

- Try different model architectures
- Experiment with reward function designs
- Train on other OpenEnv environments
- Share your trained models on Hugging Face Hub!

## Related Resources

- [OpenEnv Getting Started](../auto_getting_started/index)
- [Building Custom Environments](../auto_getting_started/plot_03_building_environments)
- [GRPO Documentation](https://huggingface.co/docs/trl/grpo_trainer)
- [Unsloth Documentation](https://github.com/unslothai/unsloth)