learning-from-experience

第1章 · Agent 基础知识 · 配套项目 chapter1/learning-from-experience

项目说明

Learning from Experience: RL vs LLM In-Context Learning

This experiment compares traditional Reinforcement Learning (Q-learning) with LLM-based in-context learning, replicating the key insights from Shunyu Yao's blog post "The Second Half".

🎯 Experiment Overview

The experiment demonstrates how LLMs can generalize through reasoning while traditional RL methods require extensive training to learn game mechanics. We use a text-based treasure hunt game with hidden mechanics that agents must discover through experience.

Key Insights Being Tested

1. Sample Efficiency: LLMs can learn from far fewer examples than traditional RL
2. Generalization: LLMs use reasoning to understand patterns, while RL memorizes state-action mappings
3. Prior Knowledge: Language pre-training provides powerful priors for reasoning about new tasks
4. Hidden Mechanics Discovery: LLMs can form hypotheses and test them, while RL requires exhaustive exploration

What You'll See

When running the LLM experiment, you'll see the complete decision-making process:

============================================================
LLM DECISION PROCESS
============================================================
📊 Experiences in memory: 15
🎮 Current room: hallway
🎯 Available actions: 8

💡 Recent successful patterns learned:
   • take red key → +5.0 reward
   • try crafting → +10.0 reward

🤔 LLM is thinking...

📝 LLM Reasoning:
----------------------------------------
  Based on my past experiences, I've learned that:
  1. The red key opens the locked door to the guard room
  2. Crafting rusty sword + magic crystal creates a silver sword
  3. The silver sword can defeat the strong guard

  Since I have the silver sword and I'm in the hallway...
----------------------------------------

✅ Chosen action: go north

This transparency shows exactly how the LLM learns and reasons, unlike the black-box nature of Q-learning!

🎮 The Game

A text-based treasure hunt game where agents must: - Navigate through multiple rooms - Collect items and keys - Defeat guards using appropriate weapons - Discover hidden mechanics through experience

Hidden Mechanics (Not Revealed to Agents)

1. Color-coded locks: Specific colored keys open matching doors
2. Weapon effectiveness: Different weapons work against different enemies
3. Crafting system: Certain items combine to create better items
4. Potion effects: Temporary abilities from consuming potions

🚀 Quick Start

Installation

# Navigate to the project directory
cd chapter1/learning-from-experience

# Install dependencies
pip install -r requirements.txt

对应书中实验 7-1（Q-learning 在寻宝游戏中的表现）与实验 7-2（传统 RL 与 LLM Agent 的对比研究）。 Q-learning 部分完全离线运行、无需任何 API Key；LLM 部分需要 Moonshot/Kimi API Key。

Setting up Kimi K3 API

To run the LLM experiments, you need a Kimi (Moonshot) API key:

1. Get your API key from Moonshot AI
2. Set the environment variable:

export MOONSHOT_API_KEY="your-api-key-here"

Or create a .env file:

echo "MOONSHOT_API_KEY=your-api-key-here" > .env

通用兜底（OpenRouter）：若未设置 MOONSHOT_API_KEY 但设置了 OPENROUTER_API_KEY， LLM 部分会自动改走 OpenRouter。由于 Kimi 模型在 OpenRouter 上不稳定可用，兜底时会使用 OPENROUTER_MODEL（默认 openai/gpt-5.6-luna）：

export OPENROUTER_API_KEY=sk-or-v1-your-key-here
python quick_demo.py   # MOONSHOT_API_KEY 缺失时自动经 OpenRouter 运行

Running the Experiment

Quick Demo (See LLM Learning in Action)

python quick_demo.py

This shows a detailed view of how the LLM learns through reasoning, displaying: - Complete thought process for each decision - How experiences accumulate and influence future decisions - The dramatic difference in learning speed vs traditional RL

Command-Line Interface (experiment.py)

experiment.py 现在提供带中文帮助的完整 CLI。查看所有参数：

python experiment.py --help

主要参数：

参数	说明	默认值
--mode {both,qlearning,rl,llm}	运行哪种智能体：qlearning/rl 只跑 Q-learning（离线）、llm 只跑 LLM Agent、both 两者对比	both
--rl-episodes	Q-learning 训练局数（实验 7-1 用 10000）	10000
--llm-episodes	LLM Agent 训练局数	20
--eval-episodes	Q-learning 训练后贪婪评估局数	100
--checkpoint-interval	学习曲线采样间隔（每 N 局记录一次胜率/Q 表规模）	1000
--model	LLM 模型名（也可用 MOONSHOT_MODEL 环境变量）	kimi-k3
--output	结果输出目录	results
--seed	随机种子，用于复现 Q-learning 学习曲线	不固定
--learning-rate / --discount / --epsilon-decay / --epsilon-min	Q-learning 超参数	0.2 / 0.99 / 0.9995 / 0.1
--stochastic	使用随机环境	确定性
--skip-llm	兼容旧用法，等价于 --mode qlearning	—

Q-Learning Only (实验 7-1，离线，无需 API)

python experiment.py --mode qlearning --rl-episodes 10000 --seed 42

训练不到 3 秒即可完成，并打印学习曲线表格，直观展现智能体如何在近万局试错中从 0% 胜率逐步学会通关（见下方"Experiment Results"）。

Full Comparison (RL vs LLM，实验 7-2)

python experiment.py --mode both --model kimi-k3

This will: 1. Train a Q-learning agent for 10000 episodes (takes ~3 seconds) and print its learning curve 2. Train an LLM agent for 20 episodes with detailed reasoning display 3. Evaluate both agents 4. Generate comparison plots 5. Save results to the results/ directory

Note: The LLM training shows the complete reasoning process for the first 3 episodes and the last episode, so you can see exactly how it learns! Each LLM episode takes ~1-2 minutes (real API calls), matching the cost trade-off discussed in 实验 7-2.

LLM Only

python experiment.py --mode llm --llm-episodes 20

Interactive Game Play

Test the game manually:

from game_environment import TreasureHuntGame

game = TreasureHuntGame()
print(game.get_state_description())
print("Available actions:", game.get_available_actions())

# Try an action
feedback, reward, done = game.execute_action("take rusty sword")
print(f"Feedback: {feedback}")
print(f"Reward: {reward}")

📊 Experiment Results

The experiment generates several outputs:

Metrics Compared

1. Sample Efficiency
2. Episodes needed to achieve good performance

3. Learning speed comparison

4. Performance

5. Victory rate in evaluation

6. Average rewards and episode lengths

7. Computational Cost

8. Training time
9. Memory usage (Q-table size vs. experience storage)
10. API calls for LLM

Visualizations

The experiment creates comparison plots showing: - Learning curves over time - Victory rate progression - Sample efficiency comparison - Key insights summary

Expected Results

Q-Learning learning curve (measured locally, --mode qlearning --rl-episodes 10000 --seed 42)

实测学习曲线（确定性环境，胜率按最近 1000 局滑动窗口统计，整段训练约 3 秒）：

Episodes	Victory rate	Q-table states	epsilon
1000	0.1%	124	0.606
2000	0.0%	129	0.368
3000	0.0%	130	0.223
5000	0.0%	132	0.100
7000	0.0%	138	0.100
8000	0.0%	140	0.100
9000	9.3%	143	0.100
10000	99.8%	143	0.100

训练后贪婪评估（--eval-episodes 100）胜率为 100%，平均 15 步通关。这正是实验 7-1 的核心现象：前 8000 局几乎 0% 胜率、只在盲目探索，价值信号需要近万局试错才传播到位。 （因随机探索有方差，未固定 --seed 时各次运行的拐点会略有不同，但整体形态一致。）

RL vs LLM（实验 7-2 的对比结论）

• Q-Learning: 需要近 10000 局才达到稳定通关；把"门/钥匙/剑"当作无意义符号，只能靠统计式暴力探索。
• LLM In-Context: 携带预训练先验，往往第一局就能在十几步内通关；靠推理理解游戏概念结构。
• Sample Efficiency: LLM 的样本效率高出 2-3 个数量级；但单局推理慢（API 调用 ~1-2 分钟）， Q-learning 跑 10000 局只需约 3 秒——权衡取决于交互成本，详见书中实验 7-2。

🏗️ Project Structure

learning-from-experience/
├── game_environment.py    # Text-based game with hidden mechanics
├── rl_agent.py            # Q-learning implementation
├── llm_agent.py           # LLM with in-context learning
├── experiment.py          # Main experiment runner
├── requirements.txt       # Python dependencies
├── README.md             # This file
└── results/              # Experiment outputs (created on run)
    └── [timestamp]/
        ├── rl_agent.pkl           # Trained Q-learning agent
        ├── llm_experiences.json   # LLM's collected experiences
        ├── experiment_results.json # Numerical results
        └── comparison_plots.png   # Visualization

🔬 Technical Details

Q-Learning Agent

• Algorithm: Tabular Q-learning with ε-greedy exploration
• State Representation: Hashed combination of room, inventory, and game state
• Learning Rate: 0.2 (configurable via --learning-rate)
• Discount Factor: 0.99 (configurable via --discount)
• Exploration: ε starts at 1.0, decays by --epsilon-decay (0.9995) to --epsilon-min (0.1)

LLM Agent (Kimi K3)

• Model: kimi-k3 (Kimi K3; override with --model or MOONSHOT_MODEL)
• Reasoning model: Kimi K3 emits a chain-of-thought (message.reasoning_content) before its final answer (message.content), so the code uses a generous max_tokens=2048 to make sure the ACTION: line is not truncated by the reasoning budget.
• Learning Method: In-context learning with experience memory (up to 50 experiences)
• Context Management: Stores successful and failed experiences
• Reasoning: Prompts LLM to reason about past experiences before acting
• Temperature: requested 0.7, but reasoning models (Kimi K3, GPT-5) only accept temperature=1, so the code auto-forces 1 for those (see _reasoning_safe_temperature)

📈 Extending the Experiment

Ideas for Further Research

1. Different Games: Try other hidden-mechanic games
2. Hybrid Approaches: Combine RL with LLM guidance
3. Transfer Learning: Test how well agents transfer to similar games
4. Ablation Studies: Remove reasoning prompts to isolate their impact
5. Other LLMs: Compare different language models

Modifying the Game

Edit game_environment.py to: - Add new rooms and items - Create more complex hidden mechanics - Adjust difficulty and rewards - Add new types of puzzles

🎓 Educational Value

This experiment demonstrates:

1. The Power of Priors: How language pre-training provides useful knowledge
2. Reasoning vs. Memorization: Different approaches to learning
3. Sample Efficiency: Why it matters for real-world applications
4. The Second Half Thesis: Moving from "can we solve it?" to "how efficiently?"

📚 References

• The Second Half by Shunyu Yao
• ReAct: Synergizing Reasoning and Acting in Language Models
• Original Q-learning paper: Watkins & Dayan (1992)

🤝 Contributing

Feel free to: - Add new hidden mechanics to the game - Implement other RL algorithms (DQN, PPO, etc.) - Try different LLM providers - Create more sophisticated evaluation metrics

📝 License

This project is for educational purposes, inspired by academic research on AI and reinforcement learning.

源代码

demo.py

#!/usr/bin/env python3
"""
Interactive demo to play the game manually or watch agents play.
"""

import os
import sys
from pathlib import Path

from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))

from game_environment import TreasureHuntGame
from rl_agent import QLearningAgent
from llm_agent import LLMAgent


def play_manual():
    """Let the user play the game manually."""
    print("\n" + "="*60)
    print("MANUAL PLAY MODE")
    print("="*60)
    print("\nYou are playing the treasure hunt game!")
    print("Try to find the dragon's treasure by exploring and discovering hidden mechanics.")

    game = TreasureHuntGame()

    while not game.game_over:
        print("\n" + "-"*40)
        print(game.get_state_description())
        print("\nAvailable actions:")
        actions = game.get_available_actions()
        for i, action in enumerate(actions, 1):
            print(f"  {i}. {action}")

        # Get user input
        choice = input("\nEnter action number or type custom action: ").strip()

        # Parse input
        if choice.isdigit() and 1 <= int(choice) <= len(actions):
            action = actions[int(choice) - 1]
        else:
            action = choice

        # Execute action
        feedback, reward, done = game.execute_action(action)
        print(f"\nFeedback: {feedback}")
        print(f"Reward: {reward:.2f}")

    if game.victory:
        print("\n🎉 CONGRATULATIONS! You won!")
    else:
        print("\n💀 GAME OVER! Better luck next time.")

    print(f"Final score: {game.score}")


def watch_rl_agent():
    """Watch a trained RL agent play."""
    print("\n" + "="*60)
    print("WATCHING Q-LEARNING AGENT")
    print("="*60)

    # Check if trained agent exists
    agent_path = Path("results") / "rl_agent_demo.pkl"

    agent = QLearningAgent()

    if agent_path.exists():
        print("Loading pre-trained agent...")
        agent.load(agent_path)
    else:
        print("No pre-trained agent found. Training one now...")
        print("This will take a few minutes...\n")

        game = TreasureHuntGame()
        agent.train(num_episodes=2000, verbose=True)

        # Save for future use
        agent_path.parent.mkdir(exist_ok=True)
        agent.save(agent_path)

    # Watch agent play
    print("\nWatching agent play...")
    game = TreasureHuntGame()
    total_reward = 0
    steps = 0

    while not game.game_over:
        print("\n" + "-"*40)
        print(game.get_state_description())

        action = agent.choose_action(game, training=False)
        print(f"\nAgent chooses: {action}")

        feedback, reward, done = game.execute_action(action)
        print(f"Feedback: {feedback}")
        print(f"Reward: {reward:.2f}")

        total_reward += reward
        steps += 1

        input("\nPress Enter to continue...")

    if game.victory:
        print("\n🎉 Agent won!")
    else:
        print("\n💀 Agent failed.")

    print(f"Total reward: {total_reward:.2f}")
    print(f"Steps taken: {steps}")


def watch_llm_agent():
    """Watch an LLM agent play with reasoning."""
    print("\n" + "="*60)
    print("WATCHING LLM AGENT (with reasoning)")
    print("="*60)

    # Check API key
    api_key = os.getenv("MOONSHOT_API_KEY")
    if not api_key and not os.getenv("OPENROUTER_API_KEY"):
        print("\nError: MOONSHOT_API_KEY not set.")
        print("Please set your Kimi API key:")
        print("  export MOONSHOT_API_KEY='your-key-here'")
        print("Or set OPENROUTER_API_KEY as a universal fallback.")
        return

    agent = LLMAgent(api_key=api_key)

    # Load experiences if available
    exp_path = Path("results") / "llm_experiences_demo.json"
    if exp_path.exists():
        print("Loading previous experiences...")
        agent.load_experiences(exp_path)
        print(f"Loaded {len(agent.experiences)} experiences")

    # Play one episode with verbose output
    print("\nWatching LLM agent play with reasoning...")
    print("(The agent will explain its thought process)\n")

    game = TreasureHuntGame()
    reward, steps, victory = agent.play_episode(game, verbose=True)

    if victory:
        print("\n🎉 LLM agent won!")
    else:
        print("\n💀 LLM agent failed.")

    print(f"Total reward: {reward:.2f}")
    print(f"Steps taken: {steps}")
    print(f"API calls made: {agent.api_calls}")

    # Save experiences
    exp_path.parent.mkdir(exist_ok=True)
    agent.save_experiences(exp_path)


def show_hidden_rules():
    """Reveal the hidden game mechanics."""
    print("\n" + "="*60)
    print("HIDDEN GAME MECHANICS (SPOILERS!)")
    print("="*60)

    game = TreasureHuntGame()
    print(game.get_hidden_rules())

    print("\nThese are the rules that agents must discover through experience.")
    print("Traditional RL requires thousands of episodes to learn these patterns,")
    print("while LLMs can often figure them out in just 20-30 episodes through reasoning.")


def main():
    """Main menu for the demo."""
    while True:
        print("\n" + "="*70)
        print("LEARNING FROM EXPERIENCE DEMO")
        print("Comparing RL vs LLM In-Context Learning")
        print("="*70)

        print("\nChoose an option:")
        print("1. Play the game manually")
        print("2. Watch Q-Learning agent play (pre-trained)")
        print("3. Watch LLM agent play with reasoning")
        print("4. Show hidden game mechanics (spoilers!)")
        print("5. Run full experiment")
        print("6. Exit")

        choice = input("\nEnter your choice (1-6): ").strip()

        if choice == "1":
            play_manual()
        elif choice == "2":
            watch_rl_agent()
        elif choice == "3":
            watch_llm_agent()
        elif choice == "4":
            show_hidden_rules()
        elif choice == "5":
            print("\nRunning full experiment...")
            os.system("python experiment.py")
        elif choice == "6":
            print("\nGoodbye!")
            break
        else:
            print("\nInvalid choice. Please try again.")


if __name__ == "__main__":
    main()

demo_stochastic.py

experiment.py

"""
Experiment runner to compare traditional RL vs LLM-based in-context learning.
This replicates the key insights from "The Second Half" blog post.
"""

import os
import json
import time
import random
import argparse
from datetime import datetime
from typing import Dict, Any, List
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

from game_environment import TreasureHuntGame
from rl_agent import QLearningAgent
from llm_agent import LLMAgent


class ExperimentRunner:
    """
    Runs experiments comparing different learning approaches.
    """

    def __init__(self, results_dir: str = "results"):
        """Initialize experiment runner."""
        self.results_dir = Path(results_dir)
        self.results_dir.mkdir(exist_ok=True)

        # Create timestamp for this experiment run
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.experiment_dir = self.results_dir / self.timestamp
        self.experiment_dir.mkdir(exist_ok=True)

        self.results = {}

    def run_rl_experiment(self,
                         num_training_episodes: int = 10000,
                         num_eval_episodes: int = 100,
                         verbose: bool = True,
                         stochastic: bool = False,
                         learning_rate: float = 0.2,
                         discount_factor: float = 0.99,
                         epsilon_decay: float = 0.9995,
                         epsilon_min: float = 0.1,
                         checkpoint_interval: int = 1000) -> Dict[str, Any]:
        """
        Run experiment with traditional Q-learning agent.

        Args:
            num_training_episodes: Number of episodes to train
            num_eval_episodes: Number of episodes to evaluate
            verbose: Whether to print details
            stochastic: Whether to use stochastic environment
            learning_rate: Q-learning learning rate (alpha)
            discount_factor: Discount factor (gamma)
            epsilon_decay: Per-episode epsilon decay
            epsilon_min: Minimum exploration rate
            checkpoint_interval: Record/print a learning-curve row every N episodes
        """
        print("\n" + "="*60)
        print("TRADITIONAL RL EXPERIMENT (Q-Learning)")
        print("="*60)

        # Initialize agent with the given hyperparameters
        agent = QLearningAgent(
            learning_rate=learning_rate,
            discount_factor=discount_factor,
            epsilon=1.0,
            epsilon_decay=epsilon_decay,
            epsilon_min=epsilon_min
        )

        # Training phase
        print(f"\nTraining for {num_training_episodes} episodes...")
        start_time = time.time()

        train_results = agent.train(
            num_episodes=num_training_episodes,
            verbose=verbose,
            stochastic=stochastic,
            checkpoint_interval=checkpoint_interval
        )

        training_time = time.time() - start_time

        # Show the learning curve (success rate over episodes) -- this is the
        # core point of experiment 7-1: the agent slowly LEARNS from experience.
        self._print_learning_curve(train_results.get("learning_curve", []),
                                    checkpoint_interval)

        # Evaluation phase
        print(f"\nEvaluating on {num_eval_episodes} episodes...")
        eval_results = agent.evaluate(
            num_episodes=num_eval_episodes,
            verbose=False,
            stochastic=stochastic
        )

        # Compile results
        results = {
            "method": "Q-Learning",
            "training_episodes": num_training_episodes,
            "training_time": training_time,
            "q_table_size": train_results["q_table_size"],
            "training_victories": train_results["total_victories"],
            "training_victory_rate": train_results["victory_rate"],
            "eval_victories": eval_results["victories"],
            "eval_victory_rate": eval_results["victory_rate"],
            "eval_avg_reward": eval_results["avg_reward"],
            "eval_avg_steps": eval_results["avg_length"],
            "episode_rewards": train_results["episode_rewards"],
            "episode_lengths": train_results["episode_lengths"],
            "learning_curve": train_results.get("learning_curve", [])
        }

        # Save agent
        agent.save(self.experiment_dir / "rl_agent.pkl")

        print(f"\nRL Training Summary:")
        print(f"  Training time: {training_time:.2f} seconds")
        print(f"  Q-table size: {train_results['q_table_size']} states")
        print(f"  Training victory rate: {train_results['victory_rate']:.2%}")
        print(f"  Evaluation victory rate: {eval_results['victory_rate']:.2%}")

        return results

    def _print_learning_curve(self, learning_curve: List[Dict[str, Any]],
                              checkpoint_interval: int):
        """Print the Q-learning success-rate-over-episodes table.

        This is the whole point of experiment 7-1: watch the victory rate climb
        from 0% (blind exploration) to ~100% only after thousands of episodes.
        """
        if not learning_curve:
            return

        window = checkpoint_interval if checkpoint_interval > 0 else 1000
        print("\n" + "-"*60)
        print(f"LEARNING CURVE (Q-Learning success rate over episodes)")
        print(f"胜率按最近 {window} 局的滑动窗口统计")
        print("-"*60)
        print(f"{'Episodes':>10} | {'Victory rate':>12} | {'Q-table':>8} | {'epsilon':>8}")
        print(f"{'-'*10}-+-{'-'*12}-+-{'-'*8}-+-{'-'*8}")
        for row in learning_curve:
            print(f"{row['episode']:>10} | {row['victory_rate']*100:>11.1f}% | "
                  f"{row['q_table_size']:>8} | {row['epsilon']:>8.3f}")
        print("-"*60)

    def run_llm_experiment(self,
                          num_training_episodes: int = 20,
                          num_eval_episodes: int = 10,
                          verbose: bool = True,
                          stochastic: bool = False,
                          model: str = "kimi-k3") -> Dict[str, Any]:
        """
        Run experiment with LLM-based in-context learning agent.

        Args:
            num_training_episodes: Number of episodes to train
            num_eval_episodes: Number of episodes to evaluate
            verbose: Whether to print details
            stochastic: Whether to use stochastic environment
        """
        print("\n" + "="*70)
        print(f"LLM-BASED IN-CONTEXT LEARNING EXPERIMENT ({model})")
        print("="*70)

        # Check for API key
        api_key = os.getenv("MOONSHOT_API_KEY")
        if not api_key and not os.getenv("OPENROUTER_API_KEY"):
            print("\n⚠️ Warning: MOONSHOT_API_KEY not set. Skipping LLM experiment.")
            print("📝 Please set your Kimi API key: export MOONSHOT_API_KEY='your-key-here'")
            print("🔗 Get your key at: https://platform.moonshot.cn/")
            print("💡 Or set OPENROUTER_API_KEY as a universal fallback.")
            return None

        print("\n✅ API key found. Initializing LLM agent...")
        print(f"🧠 Using {model} model for reasoning and in-context learning")
        print("📖 The LLM will show its complete thought process for each decision")

        # Initialize agent
        agent = LLMAgent(
            api_key=api_key,
            model=model,
            temperature=0.7,
            max_experiences=50
        )

        # Training phase (experience collection)
        print(f"\n🎓 Training Phase: Playing {num_training_episodes} episodes")
        print("💡 Watch how the LLM learns from experience without any parameter updates!")
        print("-"*70)

        start_time = time.time()

        train_results = agent.train(
            num_episodes=num_training_episodes,
            verbose=verbose,
            stochastic=stochastic
        )

        training_time = time.time() - start_time

        # Evaluation phase
        print(f"\nEvaluating on {num_eval_episodes} episodes...")
        eval_results = agent.evaluate(
            num_episodes=num_eval_episodes,
            verbose=False,
            stochastic=stochastic
        )

        # Compile results
        results = {
            "method": "LLM In-Context Learning",
            "training_episodes": num_training_episodes,
            "training_time": training_time,
            "experiences_collected": train_results["experiences_collected"],
            "api_calls": train_results["total_api_calls"],
            "total_tokens": train_results["total_tokens"],
            "training_victories": train_results["total_victories"],
            "training_victory_rate": train_results["victory_rate"],
            "eval_victories": eval_results["victories"],
            "eval_victory_rate": eval_results["victory_rate"],
            "eval_avg_reward": eval_results["avg_reward"],
            "eval_avg_steps": eval_results["avg_length"],
            "episode_rewards": train_results["episode_rewards"],
            "episode_lengths": train_results["episode_lengths"]
        }

        # Save experiences
        agent.save_experiences(self.experiment_dir / "llm_experiences.json")

        print(f"\nLLM Training Summary:")
        print(f"  Training time: {training_time:.2f} seconds")
        print(f"  Experiences collected: {train_results['experiences_collected']}")
        print(f"  API calls: {train_results['total_api_calls']}")
        print(f"  Training victory rate: {train_results['victory_rate']:.2%}")
        print(f"  Evaluation victory rate: {eval_results['victory_rate']:.2%}")

        return results

    def compare_learning_curves(self, rl_results: Dict, llm_results: Dict):
        """
        Create visualization comparing learning curves of both methods.
        """
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))

        # Plot 1: Victory rate over episodes
        ax = axes[0, 0]

        # RL victory rate (computed over windows)
        rl_rewards = rl_results["episode_rewards"]
        window_size = 100
        rl_victories = []
        for i in range(0, len(rl_rewards), window_size):
            window = rl_rewards[i:i+window_size]
            victories = sum(1 for r in window if r > 50) / len(window)
            rl_victories.append(victories)

        ax.plot(range(0, len(rl_rewards), window_size), rl_victories, 
                label=f"Q-Learning ({len(rl_rewards)} episodes)", linewidth=2)

        # LLM victory rate (per episode)
        if llm_results:
            llm_rewards = llm_results["episode_rewards"]
            llm_victories = [1 if r > 50 else 0 for r in llm_rewards]
            llm_cumulative = np.cumsum(llm_victories) / (np.arange(len(llm_victories)) + 1)
            ax.plot(range(len(llm_cumulative)), llm_cumulative,
                   label=f"LLM In-Context ({len(llm_rewards)} episodes)", linewidth=2)

        ax.set_xlabel("Episodes")
        ax.set_ylabel("Victory Rate")
        ax.set_title("Learning Progress: Victory Rate Over Time")
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Plot 2: Average reward over episodes
        ax = axes[0, 1]

        # RL rewards (smoothed)
        rl_smooth = []
        for i in range(0, len(rl_rewards), window_size):
            window = rl_rewards[i:i+window_size]
            rl_smooth.append(np.mean(window))

        ax.plot(range(0, len(rl_rewards), window_size), rl_smooth,
                label="Q-Learning", linewidth=2)

        # LLM rewards
        if llm_results:
            llm_rewards = llm_results["episode_rewards"]
            ax.plot(range(len(llm_rewards)), llm_rewards,
                   label="LLM In-Context", linewidth=2, alpha=0.7)

        ax.set_xlabel("Episodes")
        ax.set_ylabel("Episode Reward")
        ax.set_title("Learning Progress: Reward Over Time")
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Plot 3: Sample efficiency comparison
        ax = axes[1, 0]

        categories = ["Training\nEpisodes", "Evaluation\nVictory Rate", "Training\nTime (s)"]
        rl_values = [
            rl_results["training_episodes"],
            rl_results["eval_victory_rate"] * 100,
            rl_results["training_time"]
        ]

        if llm_results:
            llm_values = [
                llm_results["training_episodes"],
                llm_results["eval_victory_rate"] * 100,
                llm_results["training_time"]
            ]
        else:
            llm_values = [0, 0, 0]

        x = np.arange(len(categories))
        width = 0.35

        bars1 = ax.bar(x - width/2, rl_values, width, label='Q-Learning')
        bars2 = ax.bar(x + width/2, llm_values, width, label='LLM In-Context')

        ax.set_ylabel('Value')
        ax.set_title('Sample Efficiency Comparison')
        ax.set_xticks(x)
        ax.set_xticklabels(categories)
        ax.legend()

        # Add value labels on bars
        for bars in [bars1, bars2]:
            for bar in bars:
                height = bar.get_height()
                ax.annotate(f'{height:.1f}',
                           xy=(bar.get_x() + bar.get_width() / 2, height),
                           xytext=(0, 3),
                           textcoords="offset points",
                           ha='center', va='bottom')

        # Plot 4: Key insights text
        ax = axes[1, 1]
        ax.axis('off')

        insights = [
            "KEY INSIGHTS (Replicating 'The Second Half' Findings):",
            "",
            "1. SAMPLE EFFICIENCY:",
            f"   • Q-Learning: {rl_results['training_episodes']} episodes needed",
            f"   • LLM: {llm_results['training_episodes'] if llm_results else 'N/A'} episodes needed",
            f"   • Improvement: {rl_results['training_episodes'] / (llm_results['training_episodes'] if llm_results and llm_results['training_episodes'] > 0 else 1):.1f}x fewer samples",
            "",
            "2. GENERALIZATION:",
            "   • Q-Learning: Learns specific state-action mappings",
            "   • LLM: Reasons about patterns and transfers knowledge",
            "",
            "3. HIDDEN MECHANICS DISCOVERY:",
            "   • Q-Learning: Requires extensive exploration",
            "   • LLM: Can hypothesize and test theories",
            "",
            "4. COMPUTATIONAL TRADE-OFF:",
            f"   • Q-Learning: Fast inference, slow learning",
            f"   • LLM: Slower inference (API calls), fast adaptation"
        ]

        y_pos = 0.9
        for line in insights:
            if line.startswith("KEY INSIGHTS"):
                ax.text(0.5, y_pos, line, transform=ax.transAxes,
                       fontsize=12, fontweight='bold', ha='center')
            elif line.startswith(("1.", "2.", "3.", "4.")):
                ax.text(0.1, y_pos, line, transform=ax.transAxes,
                       fontsize=11, fontweight='bold')
            else:
                ax.text(0.1, y_pos, line, transform=ax.transAxes,
                       fontsize=10)
            y_pos -= 0.06

        plt.tight_layout()
        plt.savefig(self.experiment_dir / "comparison_plots.png", dpi=150)
        plt.show()

        print(f"\nPlots saved to {self.experiment_dir / 'comparison_plots.png'}")

    def run_full_experiment(self,
                           rl_episodes: int = 10000,
                           llm_episodes: int = 20,
                           eval_episodes: int = 100,
                           verbose: bool = False,
                           stochastic: bool = False,
                           model: str = "kimi-k3",
                           checkpoint_interval: int = 1000,
                           rl_hyperparams: Dict[str, float] = None):
        """
        Run full comparison experiment.

        Args:
            rl_episodes: Number of episodes for RL training
            llm_episodes: Number of episodes for LLM training
            eval_episodes: Number of episodes for RL evaluation
            verbose: Whether to print details
            stochastic: Whether to use stochastic environment
            model: LLM model name (Moonshot/Kimi)
            checkpoint_interval: Learning-curve sampling interval (RL)
            rl_hyperparams: Optional dict overriding Q-learning hyperparameters
        """
        print("\n" + "="*70)
        print("EXPERIMENT: Traditional RL vs LLM In-Context Learning")
        print("Replicating insights from 'The Second Half' by Shunyu Yao")
        print("="*70)

        # Show game rules for reference
        game = TreasureHuntGame(stochastic=stochastic)
        print("\n" + game.get_hidden_rules())

        rl_hyperparams = rl_hyperparams or {}

        # Run RL experiment
        rl_results = self.run_rl_experiment(
            num_training_episodes=rl_episodes,
            num_eval_episodes=eval_episodes,
            verbose=verbose,
            stochastic=stochastic,
            checkpoint_interval=checkpoint_interval,
            **rl_hyperparams
        )
        self.results["rl"] = rl_results

        # Run LLM experiment
        llm_results = self.run_llm_experiment(
            num_training_episodes=llm_episodes,
            num_eval_episodes=10,
            verbose=verbose,
            stochastic=stochastic,
            model=model
        )
        self.results["llm"] = llm_results

        # Save combined results
        with open(self.experiment_dir / "experiment_results.json", 'w') as f:
            json.dump(self.results, f, indent=2)

        # Generate comparison plots
        if llm_results:
            self.compare_learning_curves(rl_results, llm_results)

        # Print final comparison
        print("\n" + "="*70)
        print("EXPERIMENT RESULTS SUMMARY")
        print("="*70)

        print("\n1. SAMPLE EFFICIENCY:")
        print(f"   Q-Learning needed {rl_results['training_episodes']} episodes")
        if llm_results:
            print(f"   LLM needed {llm_results['training_episodes']} episodes")
            print(f"   → LLM is {rl_results['training_episodes'] / llm_results['training_episodes']:.1f}x more sample efficient")

        print("\n2. PERFORMANCE:")
        print(f"   Q-Learning eval victory rate: {rl_results['eval_victory_rate']:.2%}")
        if llm_results:
            print(f"   LLM eval victory rate: {llm_results['eval_victory_rate']:.2%}")

        print("\n3. COMPUTATIONAL COST:")
        print(f"   Q-Learning: {rl_results['training_time']:.2f} seconds, {rl_results['q_table_size']} states")
        if llm_results:
            print(f"   LLM: {llm_results['training_time']:.2f} seconds, {llm_results['api_calls']} API calls")

        print(f"\nResults saved to: {self.experiment_dir}")

        return self.results


def main():
    """Main entry point for the experiment."""
    parser = argparse.ArgumentParser(
        description="实验 7-1 / 7-2：在寻宝游戏中对比 Q-learning 与 LLM 的\"从经验中学习\"。"
                    "Q-learning 完全离线运行（无需 API），LLM 模式需要 Moonshot/Kimi API Key。",
        epilog="示例：\n"
               "  python experiment.py --mode qlearning              # 只跑 Q-learning（离线，输出学习曲线）\n"
               "  python experiment.py --mode qlearning --rl-episodes 10000 --seed 42\n"
               "  python experiment.py --mode both --model kimi-k3  # RL vs LLM 对比\n"
               "  python experiment.py --mode llm --llm-episodes 20   # 只跑 LLM 智能体",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--mode", choices=["both", "qlearning", "rl", "llm"], default="both",
        help="运行哪种智能体：qlearning/rl 只跑 Q-learning（离线），"
             "llm 只跑 LLM 智能体（需 API），both 两者对比（默认）"
    )
    parser.add_argument(
        "--rl-episodes", type=int, default=10000,
        help="Q-learning 训练局数（默认 10000，对应实验 7-1）"
    )
    parser.add_argument(
        "--llm-episodes", type=int, default=20,
        help="LLM 智能体的训练局数（默认 20）"
    )
    parser.add_argument(
        "--eval-episodes", type=int, default=100,
        help="Q-learning 训练后贪婪评估的局数（默认 100）"
    )
    parser.add_argument(
        "--checkpoint-interval", type=int, default=1000,
        help="学习曲线采样间隔：每 N 局记录一次胜率/Q 表规模（默认 1000）"
    )
    parser.add_argument(
        "--model", type=str, default=os.getenv("MOONSHOT_MODEL", "kimi-k3"),
        help="LLM 模型名称（Moonshot/Kimi，默认 kimi-k3，可用 MOONSHOT_MODEL 环境变量覆盖）"
    )
    parser.add_argument(
        "--output", type=str, default="results",
        help="结果输出目录（默认 results/，每次运行会新建时间戳子目录）"
    )
    parser.add_argument(
        "--seed", type=int, default=None,
        help="随机种子，用于复现 Q-learning 的学习曲线（默认不固定）"
    )
    # Q-learning 超参数
    parser.add_argument("--learning-rate", type=float, default=0.2,
                        help="Q-learning 学习率 alpha（默认 0.2）")
    parser.add_argument("--discount", type=float, default=0.99,
                        help="折扣因子 gamma（默认 0.99）")
    parser.add_argument("--epsilon-decay", type=float, default=0.9995,
                        help="每局 epsilon 衰减系数（默认 0.9995）")
    parser.add_argument("--epsilon-min", type=float, default=0.1,
                        help="最小探索率 epsilon（默认 0.1）")
    parser.add_argument(
        "--verbose", action="store_true",
        help="训练过程中打印详细信息"
    )
    parser.add_argument(
        "--skip-llm", action="store_true",
        help="[兼容旧用法] 跳过 LLM 实验，等价于 --mode qlearning"
    )
    parser.add_argument(
        "--stochastic", action="store_true",
        help="使用随机环境（奖励与动作带随机扰动）"
    )
    parser.add_argument(
        "--deterministic", action="store_true",
        help="使用确定性环境（默认）"
    )

    args = parser.parse_args()

    # Handle environment mode
    if args.deterministic and args.stochastic:
        print("Error: Cannot specify both --deterministic and --stochastic")
        return

    # Episode counts must be positive; 0 would divide by zero in the
    # train/evaluate victory-rate and average calculations.
    if args.rl_episodes < 1 or args.llm_episodes < 1 or args.eval_episodes < 1:
        print("Error: --rl-episodes, --llm-episodes and --eval-episodes must all be >= 1")
        return

    # Resolve run mode (--skip-llm kept as a backwards-compatible alias)
    mode = "qlearning" if args.skip_llm else args.mode

    # Seed for reproducible Q-learning learning curves
    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)
        print(f"\n🎲 Random seed set to {args.seed} for reproducibility")

    stochastic = args.stochastic  # Default is False (deterministic)

    if stochastic:
        print("\n🎲 Running experiment with STOCHASTIC environment")
        print("  - Random reward variations")
        print("  - 3% chance of action failure")
        print("  - Combat and crafting variations\n")
    else:
        print("\n🎯 Running experiment with DETERMINISTIC environment\n")

    rl_hyperparams = {
        "learning_rate": args.learning_rate,
        "discount_factor": args.discount,
        "epsilon_decay": args.epsilon_decay,
        "epsilon_min": args.epsilon_min,
    }

    # Run experiment
    runner = ExperimentRunner(results_dir=args.output)

    if mode in ("qlearning", "rl"):
        # Run only the Q-learning experiment (fully offline, no API needed)
        rl_results = runner.run_rl_experiment(
            num_training_episodes=args.rl_episodes,
            num_eval_episodes=args.eval_episodes,
            verbose=args.verbose,
            stochastic=stochastic,
            checkpoint_interval=args.checkpoint_interval,
            **rl_hyperparams
        )
        runner.results["rl"] = rl_results
        with open(runner.experiment_dir / "experiment_results.json", 'w') as f:
            json.dump(runner.results, f, indent=2)
        print(f"\nResults saved to: {runner.experiment_dir}")
        print("\nSkipped LLM experiment. Use --mode both with an API key to compare.")
    elif mode == "llm":
        # Run only the LLM experiment
        llm_results = runner.run_llm_experiment(
            num_training_episodes=args.llm_episodes,
            verbose=args.verbose,
            stochastic=stochastic,
            model=args.model
        )
        runner.results["llm"] = llm_results
        with open(runner.experiment_dir / "experiment_results.json", 'w') as f:
            json.dump(runner.results, f, indent=2)
        print(f"\nResults saved to: {runner.experiment_dir}")
    else:
        # Run full comparison
        results = runner.run_full_experiment(
            rl_episodes=args.rl_episodes,
            llm_episodes=args.llm_episodes,
            eval_episodes=args.eval_episodes,
            verbose=args.verbose,
            stochastic=stochastic,
            model=args.model,
            checkpoint_interval=args.checkpoint_interval,
            rl_hyperparams=rl_hyperparams
        )

    print("\nExperiment complete!")


if __name__ == "__main__":
    main()

game_environment.py

"""
Text-based treasure hunt game with hidden mechanics.
Inspired by Shunyu Yao's insights on reasoning and generalization in AI.
"""

import random
from typing import Dict, List, Tuple, Optional, Set
from dataclasses import dataclass, field
from enum import Enum


class ItemType(Enum):
    KEY = "key"
    WEAPON = "weapon"
    TREASURE = "treasure"
    TOOL = "tool"
    POTION = "potion"


@dataclass
class Item:
    name: str
    item_type: ItemType
    description: str
    properties: Dict[str, any] = field(default_factory=dict)


@dataclass
class Room:
    name: str
    description: str
    items: List[Item] = field(default_factory=list)
    exits: Dict[str, str] = field(default_factory=dict)  # direction -> room_name
    locked_exits: Dict[str, str] = field(default_factory=dict)  # direction -> required_key
    has_guard: bool = False
    guard_defeated: bool = False


class TreasureHuntGame:
    """
    A text-based game with hidden mechanics that agents must discover:
    1. Certain colored keys open corresponding colored doors
    2. Guards block access to treasures and require specific weapons
    3. Some items combine to create new items (hidden crafting)
    4. Potions provide temporary abilities
    """

    def __init__(self, seed: int = None, stochastic: bool = False):
        """
        Initialize the game environment.

        Args:
            seed: Random seed for reproducibility
            stochastic: If True, adds random elements to the game
        """
        if seed is not None:
            random.seed(seed)

        self.stochastic = stochastic
        self.random_state = random.Random(seed) if stochastic else None

        self.rooms = {}
        self.current_room = None
        self.inventory = []
        self.score = 0
        self.moves = 0
        self.max_moves = 50  # Reduced for faster episodes
        self.game_over = False
        self.victory = False
        self.active_effects = {}

        # Hidden mechanics (not revealed to agents initially)
        self.color_key_mapping = {
            "red key": "red door",
            "blue key": "blue door",
            "golden key": "golden door"
        }

        self.weapon_effectiveness = {
            "rusty sword": ["weak guard"],
            "silver sword": ["weak guard", "strong guard", "dragon"]
        }

        self.crafting_recipes = {
            frozenset(["rusty sword", "magic crystal"]): "silver sword"
        }

        self._initialize_world()

    def _initialize_world(self):
        """Create the game world with rooms and items - simplified for better learning."""

        # Create a simpler world that's easier to learn but still demonstrates the concepts
        self.rooms["entrance"] = Room(
            name="entrance",
            description="You stand in a dimly lit entrance hall. Stone walls echo your footsteps.",
            items=[
                Item("rusty sword", ItemType.WEAPON, "An old sword with rust spots")
            ],
            exits={"north": "hallway", "east": "storage"}
        )

        self.rooms["storage"] = Room(
            name="storage",
            description="A dusty storage room filled with old crates and barrels.",
            items=[
                Item("red key", ItemType.KEY, "A small red metal key"),
                Item("magic crystal", ItemType.TOOL, "A glowing crystal that hums with energy")
            ],
            exits={"west": "entrance"}
        )

        self.rooms["hallway"] = Room(
            name="hallway",
            description="A long hallway with a locked door to the north.",
            exits={"south": "entrance", "north": "guard_room"},
            locked_exits={"north": "red key"}
        )

        self.rooms["guard_room"] = Room(
            name="guard_room",
            description="A large room with weapon racks. A guard blocks the treasure!",
            has_guard=True,
            items=[],
            exits={"south": "hallway", "east": "treasure_room"}
        )

        self.rooms["treasure_room"] = Room(
            name="treasure_room",
            description="The treasure room! Gold coins and jewels sparkle in the light.",
            items=[
                Item("dragon's treasure", ItemType.TREASURE, "A massive hoard of gold and gems", 
                     {"value": 1000})
            ],
            exits={"west": "guard_room"}
        )

        self.current_room = self.rooms["entrance"]

    def get_state_description(self) -> str:
        """Get a natural language description of the current game state."""
        desc = []
        desc.append(f"\n=== Room: {self.current_room.name.replace('_', ' ').title()} ===")
        desc.append(self.current_room.description)

        if self.current_room.has_guard and not self.current_room.guard_defeated:
            desc.append("A guard blocks your way!")

        if self.current_room.items:
            desc.append("\nYou see:")
            for item in self.current_room.items:
                desc.append(f"  - {item.name}: {item.description}")

        exits = []
        for direction, room in self.current_room.exits.items():
            if direction in self.current_room.locked_exits:
                exits.append(f"{direction} (locked)")
            else:
                exits.append(direction)
        desc.append(f"\nExits: {', '.join(exits)}")

        if self.inventory:
            desc.append(f"\nInventory: {', '.join([item.name for item in self.inventory])}")

        desc.append(f"\nScore: {self.score} | Moves: {self.moves}/{self.max_moves}")

        return "\n".join(desc)

    def get_available_actions(self) -> List[str]:
        """Get list of available actions in current state."""
        actions = []

        # Movement actions
        for direction in self.current_room.exits.keys():
            actions.append(f"go {direction}")

        # Item actions
        for item in self.current_room.items:
            actions.append(f"take {item.name}")

        for item in self.inventory:
            actions.append(f"use {item.name}")
            actions.append(f"drop {item.name}")

        # Combat actions
        if self.current_room.has_guard and not self.current_room.guard_defeated:
            for item in self.inventory:
                if item.item_type == ItemType.WEAPON:
                    actions.append(f"attack with {item.name}")

        # Special actions
        actions.append("look around")
        actions.append("check inventory")

        # Crafting (if player has discovered it)
        if len(self.inventory) >= 2:
            actions.append("try crafting")

        return actions

    def execute_action(self, action: str) -> Tuple[str, float, bool]:
        """
        Execute an action and return (feedback, reward, done).
        """
        if self.game_over:
            return "Game is already over.", 0, True

        self.moves += 1
        action = action.lower().strip()

        # Base reward with stochastic variation
        if self.stochastic:
            # Add small random variation to rewards
            reward = -0.5 + self.random_state.uniform(-0.1, 0.1)

            # Small chance of action failure in stochastic mode
            if self.random_state.random() < 0.03:  # 3% chance
                return "You fumble and need to try again.", reward - 0.2, False
        else:
            reward = -0.5  # Negative reward for each move to encourage efficiency

        # Check move limit
        if self.moves >= self.max_moves:
            self.game_over = True
            return "You've run out of moves! Game over.", -10, True

        # Parse action
        if action.startswith("go "):
            direction = action[3:]
            result, move_reward = self._move(direction)
            reward += move_reward

        elif action.startswith("take "):
            item_name = action[5:]
            result, take_reward = self._take_item(item_name)
            reward += take_reward

        elif action.startswith("use "):
            item_name = action[4:]
            result, use_reward = self._use_item(item_name)
            reward += use_reward

        elif action.startswith("drop "):
            item_name = action[5:]
            result = self._drop_item(item_name)

        elif action.startswith("attack with "):
            weapon_name = action[12:]
            result, attack_reward = self._attack(weapon_name)
            reward += attack_reward

        elif action == "look around":
            result = self.get_state_description()

        elif action == "check inventory":
            if self.inventory:
                result = "Inventory: " + ", ".join([f"{item.name} ({item.item_type.value})" 
                                                    for item in self.inventory])
            else:
                result = "Your inventory is empty."

        elif action == "try crafting":
            result, craft_reward = self._try_crafting()
            reward += craft_reward

        else:
            result = f"Unknown action: {action}"
            reward -= 1

        # Check victory condition
        if self._check_victory():
            self.victory = True
            self.game_over = True
            reward += 100
            result += "\n\n🎉 VICTORY! You've collected the dragon's treasure!"

        return result, reward, self.game_over

    def _move(self, direction: str) -> Tuple[str, float]:
        """Move to another room."""
        if direction not in self.current_room.exits:
            return f"You can't go {direction} from here.", -1

        # Check if locked
        if direction in self.current_room.locked_exits:
            required_key = self.current_room.locked_exits[direction]
            if not any(item.name == required_key for item in self.inventory):
                return f"The {direction} exit is locked. You need a {required_key}.", -0.5
            else:
                # Unlock and move
                del self.current_room.locked_exits[direction]
                room_name = self.current_room.exits[direction]
                self.current_room = self.rooms[room_name]
                return f"You unlock the door with the {required_key} and move {direction}.", 5

        # Check for guard
        if self.current_room.has_guard and not self.current_room.guard_defeated:
            return "A guard blocks your way! You must defeat them first.", -1

        # Move to new room
        room_name = self.current_room.exits[direction]
        self.current_room = self.rooms[room_name]
        return f"You move {direction} to the {self.current_room.name}.", 1

    def _take_item(self, item_name: str) -> Tuple[str, float]:
        """Pick up an item."""
        for item in self.current_room.items:
            if item.name.lower() == item_name.lower():
                self.current_room.items.remove(item)
                self.inventory.append(item)

                # Reward based on item type
                if item.item_type == ItemType.TREASURE:
                    reward = 100  # Big reward for getting the treasure!
                elif item.item_type == ItemType.KEY:
                    reward = 5
                elif item.item_type == ItemType.WEAPON:
                    reward = 3
                else:
                    reward = 2

                # Add stochastic variation
                if self.stochastic:
                    reward += self.random_state.uniform(-0.5, 0.5)

                return f"You take the {item.name}.", reward

        penalty = -0.5
        if self.stochastic:
            penalty += self.random_state.uniform(-0.1, 0.1)

        return f"There's no {item_name} here.", penalty

    def _drop_item(self, item_name: str) -> str:
        """Drop an item."""
        for item in self.inventory:
            if item.name.lower() == item_name.lower():
                self.inventory.remove(item)
                self.current_room.items.append(item)
                return f"You drop the {item.name}."

        return f"You don't have a {item_name}."

    def _use_item(self, item_name: str) -> Tuple[str, float]:
        """Use an item."""
        for item in self.inventory:
            if item.name.lower() == item_name.lower():
                if item.item_type == ItemType.POTION:
                    self.inventory.remove(item)
                    if "healing" in item.name:
                        return "You drink the healing potion and feel refreshed!", 5
                    elif "strength" in item.name:
                        self.active_effects["strength"] = 10
                        return "You feel a surge of power! Your attacks will be stronger.", 5

                elif item.item_type == ItemType.KEY:
                    # Keys are used automatically when moving
                    return f"The {item.name} will be used automatically when needed.", 0

                else:
                    return f"You can't use the {item.name} right now.", -0.5

        return f"You don't have a {item_name}.", -0.5

    def _attack(self, weapon_name: str) -> Tuple[str, float]:
        """Attack with a weapon."""
        if not self.current_room.has_guard or self.current_room.guard_defeated:
            return "There's nothing to attack here.", -1

        weapon = None
        for item in self.inventory:
            if item.name.lower() == weapon_name.lower():
                weapon = item
                break

        if not weapon:
            return f"You don't have a {weapon_name}.", -1

        if weapon.item_type != ItemType.WEAPON:
            return f"The {weapon_name} is not a weapon!", -1

        # Check weapon effectiveness (hidden mechanic)
        # In our simplified game, the guard in guard_room is a "strong guard"
        guard_type = "strong guard"

        if weapon.name in self.weapon_effectiveness:
            if guard_type in self.weapon_effectiveness[weapon.name]:
                # In stochastic mode, add combat variations
                if self.stochastic:
                    roll = self.random_state.random()
                    if roll < 0.1:  # 10% critical hit
                        self.current_room.guard_defeated = True
                        return f"Critical hit! You defeat the {guard_type} with your {weapon.name}!", 30
                    elif roll < 0.95:  # 85% normal success
                        self.current_room.guard_defeated = True
                        return f"You defeat the {guard_type} with your {weapon.name}!", 20
                    else:  # 5% glancing blow
                        return f"Your attack glances off! The {guard_type} is still standing.", -0.5
                else:
                    self.current_room.guard_defeated = True
                    return f"You defeat the {guard_type} with your {weapon.name}!", 20
            else:
                penalty = -2
                if self.stochastic:
                    penalty += self.random_state.uniform(-0.5, 0.5)
                return f"Your {weapon.name} is not effective against the {guard_type}!", penalty

        return f"Your {weapon.name} doesn't seem to work.", -1

    def _try_crafting(self) -> Tuple[str, float]:
        """Try to craft items (hidden mechanic)."""
        if len(self.inventory) < 2:
            return "You need at least two items to craft.", -0.5

        # Check all possible combinations
        inventory_names = [item.name for item in self.inventory]

        for recipe, result in self.crafting_recipes.items():
            if recipe.issubset(set(inventory_names)):
                # In stochastic mode, crafting might have variations
                if self.stochastic:
                    if self.random_state.random() < 0.9:  # 90% success rate
                        # Craft the item
                        for ingredient in recipe:
                            for item in self.inventory[:]:
                                if item.name == ingredient:
                                    self.inventory.remove(item)
                                    break

                        new_item = self._create_item(result)
                        self.inventory.append(new_item)
                        reward = 10 + self.random_state.uniform(-1, 2)
                        return f"You successfully craft a {result}!", reward
                    else:
                        # 10% chance of crafting mishap (items not consumed)
                        return "The crafting attempt fizzles. Try again!", -0.2
                else:
                    # Deterministic crafting
                    for ingredient in recipe:
                        for item in self.inventory[:]:
                            if item.name == ingredient:
                                self.inventory.remove(item)
                                break

                    new_item = self._create_item(result)
                    self.inventory.append(new_item)
                    return f"You successfully craft a {result}!", 10  # Good reward for discovering crafting

        penalty = -0.5
        if self.stochastic:
            penalty += self.random_state.uniform(-0.1, 0.1)

        return "These items don't combine into anything useful.", penalty

    def _create_item(self, item_name: str) -> Item:
        """Create an item by name."""
        if item_name == "silver sword":
            return Item("silver sword", ItemType.WEAPON, "A gleaming silver blade")
        elif item_name == "magic staff":
            return Item("magic staff", ItemType.WEAPON, "A staff crackling with magical energy")
        else:
            return Item(item_name, ItemType.TOOL, "A crafted item")

    def _check_victory(self) -> bool:
        """Check if the player has won."""
        for item in self.inventory:
            if item.name == "dragon's treasure":
                return True
        return False

    def reset(self, seed: int = None) -> str:
        """Reset the game to initial state."""
        if seed is None:
            seed = random.randint(0, 10000)
        self.__init__(seed=seed, stochastic=self.stochastic)
        return self.get_state_description()

    def get_hidden_rules(self) -> str:
        """Return the hidden game rules (for debugging/analysis)."""
        rules = []
        rules.append("Hidden Game Mechanics (Simplified Version):")
        rules.append("\n1. To win the game:")
        rules.append("   - Get the red key from storage room")
        rules.append("   - Use it to unlock the door to guard room")
        rules.append("   - Craft a silver sword (rusty sword + magic crystal)")
        rules.append("   - Defeat the strong guard with the silver sword")
        rules.append("   - Collect the dragon's treasure")

        rules.append("\n2. Key mechanics:")
        rules.append("   - Red key opens the locked door in the hallway")

        rules.append("\n3. Weapon effectiveness:")
        for weapon, targets in self.weapon_effectiveness.items():
            rules.append(f"   - {weapon} defeats: {', '.join(targets)}")

        rules.append("\n4. Crafting recipe:")
        for ingredients, result in self.crafting_recipes.items():
            rules.append(f"   - {' + '.join(ingredients)} = {result}")

        rules.append("\n5. Optimal solution:")
        rules.append("   - Takes about 10-15 moves if done efficiently")

        return "\n".join(rules)

llm_agent.py

"""
LLM-based Agent using In-Context Learning with the Kimi (Moonshot) API.
This demonstrates how LLMs can generalize through reasoning without extensive training.
Default model is Kimi K3 (matching 实验 7-2 in the book); override via the
`model` argument or the MOONSHOT_MODEL environment variable.
"""

import os
import json
import time
from typing import Dict, List, Tuple, Any, Optional
from dataclasses import dataclass, asdict
import openai
from game_environment import TreasureHuntGame


def _reasoning_safe_temperature(model, requested=1.0):
    """Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
    Return 1 for those; otherwise the requested value so non-reasoning
    providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
    m = str(model or "").lower().replace("/", "-")
    return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested


def _map_model_to_openrouter(model):
    """Map a bare model id to an OpenRouter model id.
    - ids already containing '/' -> left as-is
    - gpt-*/o1-*/o3-*/o4-* -> 'openai/<id>'
    - claude-* -> anthropic Claude (opus/sonnet/haiku)
    - other native ids (kimi-*, doubao-*, ...) are NOT reliably on OpenRouter,
      so fall back to OPENROUTER_MODEL or a safe default that always works.
    """
    m = (model or "").strip()
    if "/" in m:
        return m
    ml = m.lower()
    if ml.startswith(("gpt-", "o1-", "o3-", "o4-")):
        return "openai/" + m
    if ml.startswith("claude-"):
        if "sonnet" in ml:
            return "anthropic/claude-sonnet-4.6"
        if "haiku" in ml:
            return "anthropic/claude-haiku-4.5"
        return "anthropic/claude-opus-4.8"
    if ml.startswith("kimi"):
        # kimi-k3 is not on OpenRouter; moonshotai/kimi-k2.6 is the closest hosted id.
        return "moonshotai/kimi-k2.6"
    return os.getenv("OPENROUTER_MODEL", "openai/gpt-5.6-luna")


def resolve_llm_backend(primary_key, primary_base_url, model):
    """Universal OpenRouter fallback for LLM backend resolution.

    Returns (api_key, base_url, model, using_openrouter).
    - If the primary provider key is present, behavior is unchanged.
    - Else if OPENROUTER_API_KEY is present, route through OpenRouter and map
      the model id to an OpenRouter id.
    - Else raise a clear error listing the accepted keys.
    """
    openrouter_key = os.getenv("OPENROUTER_API_KEY")
    # gpt-5.x (incl. gpt-5.6*) needs OpenAI org-verification on the direct API;
    # when an OpenRouter key is present, prefer routing these ids through it.
    if openrouter_key and str(model or "").lower().startswith("gpt-5"):
        base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
        return openrouter_key, base_url, _map_model_to_openrouter(model), True
    if primary_key:
        return primary_key, primary_base_url, model, False
    if openrouter_key:
        base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
        return openrouter_key, base_url, _map_model_to_openrouter(model), True
    raise ValueError(
        "No API key found. Set MOONSHOT_API_KEY (primary) or "
        "OPENROUTER_API_KEY (universal fallback)."
    )


@dataclass
class GameExperience:
    """Represents a single game interaction experience."""
    state_description: str
    action: str
    feedback: str
    reward: float
    success: bool  # Whether the action led to positive outcome


class LLMAgent:
    """
    LLM-based agent that uses in-context learning to play the game.
    Stores experiences and uses them to reason about future actions.
    """

    def __init__(self,
                 api_key: str = None,
                 model: str = "kimi-k3",  # Kimi K3 (see 实验 7-2)
                 base_url: str = "https://api.moonshot.cn/v1",
                 temperature: float = 0.7,
                 max_experiences: int = 50):
        """
        Initialize LLM agent with the Kimi (Moonshot) API.

        Args:
            api_key: Kimi API key (or set MOONSHOT_API_KEY env var)
            model: Model name (Moonshot/Kimi, default kimi-k3)
            base_url: API base URL
            temperature: Sampling temperature for generation
            max_experiences: Maximum number of experiences to store
        """
        # Set up API client (Moonshot primary, OpenRouter universal fallback)
        primary_key = api_key or os.getenv("MOONSHOT_API_KEY")
        self.api_key, resolved_base_url, self.model, self.using_openrouter = \
            resolve_llm_backend(primary_key, base_url, model)
        if self.using_openrouter:
            print(f"ℹ️  MOONSHOT_API_KEY not set; routing via OpenRouter (model: {self.model})")

        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url=resolved_base_url
        )
        self.temperature = temperature

        # Experience memory for in-context learning
        self.experiences: List[GameExperience] = []
        self.max_experiences = max_experiences

        # Statistics
        self.episode_rewards = []
        self.episode_lengths = []
        self.victories = 0
        self.total_episodes = 0
        self.api_calls = 0
        self.total_tokens = 0

    def _build_context(self, current_state: str, available_actions: List[str]) -> str:
        """
        Build context for the LLM including task description and past experiences.
        This is the key to in-context learning.
        """
        context = []

        # Task description
        context.append("""You are playing a text-based treasure hunt game. Your goal is to find and collect the dragon's treasure.

The game has hidden mechanics that you need to discover through experience:
- Certain items may be required to unlock doors or defeat guards
- Items might combine to create better items
- Different weapons have different effectiveness

You should reason about what you've learned from past experiences to make better decisions.""")

        # Add relevant past experiences
        if self.experiences:
            context.append("\n=== PAST EXPERIENCES ===")
            context.append("Here are some experiences from previous attempts that might help you:")

            # Group experiences by pattern
            successful_patterns = []
            failed_patterns = []

            for exp in self.experiences[-self.max_experiences:]:
                exp_text = f"State: {exp.state_description[:200]}...\nAction: {exp.action}\nResult: {exp.feedback}\nReward: {exp.reward:.1f}"

                if exp.success:
                    successful_patterns.append(exp_text)
                else:
                    failed_patterns.append(exp_text)

            if successful_patterns:
                context.append("\n** Successful actions:")
                for pattern in successful_patterns[-10:]:  # Last 10 successful
                    context.append(pattern)

            if failed_patterns:
                context.append("\n** Failed actions to avoid:")
                for pattern in failed_patterns[-5:]:  # Last 5 failed
                    context.append(pattern)

        # Current situation
        context.append("\n=== CURRENT SITUATION ===")
        context.append(current_state)
        context.append(f"\nAvailable actions: {', '.join(available_actions)}")

        return "\n".join(context)

    def _build_prompt(self, context: str) -> str:
        """Build the full prompt for the LLM."""
        prompt = f"""{context}

Based on your understanding of the game mechanics from past experiences and the current situation, reason step-by-step about what action to take:

1. What have you learned from past experiences that applies here?
2. What is your current goal or sub-goal?
3. Which available action best helps achieve that goal?

Think through this carefully, then provide your chosen action.

IMPORTANT: Your response must end with exactly one line starting with "ACTION:" followed by one of the available actions listed above.

Example format:
[Your reasoning here...]
ACTION: take red key
"""
        return prompt

    def choose_action(self, game: TreasureHuntGame, verbose: bool = True) -> str:
        """
        Choose an action using LLM reasoning with in-context learning.
        """
        # Get current state and available actions
        state_description = game.get_state_description()
        available_actions = game.get_available_actions()

        if not available_actions:
            return "look around"

        # Build context with past experiences
        context = self._build_context(state_description, available_actions)
        prompt = self._build_prompt(context)

        if verbose:
            print("\n" + "="*60)
            print("LLM DECISION PROCESS")
            print("="*60)
            print(f"📊 Experiences in memory: {len(self.experiences)}")
            print(f"🎮 Current room: {game.current_room.name}")
            print(f"🎯 Available actions: {len(available_actions)}")

            # Show some recent successful experiences if any
            successful = [e for e in self.experiences if e.success]
            if successful:
                print(f"\n💡 Recent successful patterns learned:")
                for exp in successful[-3:]:
                    print(f"   • {exp.action} → +{exp.reward:.1f} reward")

        try:
            print("\n🤔 LLM is thinking...")

            # Call Kimi K3 API.
            # Kimi K3 is a REASONING model: its chain-of-thought is billed as
            # completion tokens (returned separately in message.reasoning_content),
            # and the final answer in message.content is only emitted AFTER the
            # reasoning finishes. A small max_tokens can therefore be consumed
            # entirely by reasoning, truncating content to empty and breaking the
            # ACTION parse. Keep a generous budget so the ACTION line survives.
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are an intelligent game-playing agent that learns from experience."},
                    {"role": "user", "content": prompt}
                ],
                temperature=_reasoning_safe_temperature(self.model, self.temperature),
                max_tokens=2048
            )

            self.api_calls += 1
            if hasattr(response.usage, 'total_tokens'):
                self.total_tokens += response.usage.total_tokens

            # Extract action from response (final answer lives in .content;
            # reasoning models expose their thinking separately in
            # .reasoning_content, which we don't need for action parsing).
            response_text = response.choices[0].message.content or ""

            if verbose:
                print("\n📝 LLM Reasoning:")
                print("-" * 40)
                # Show the reasoning part (before ACTION:)
                reasoning_lines = []
                for line in response_text.split('\n'):
                    if line.startswith("ACTION:"):
                        break
                    if line.strip():
                        reasoning_lines.append(line)

                # Show last few lines of reasoning
                for line in reasoning_lines[-5:]:
                    print(f"  {line[:100]}...")  # Truncate long lines
                print("-" * 40)

            # Parse action from response
            lines = response_text.strip().split('\n')
            for line in reversed(lines):
                if line.startswith("ACTION:"):
                    action = line[7:].strip()

                    # Validate action
                    if action in available_actions:
                        if verbose:
                            print(f"\n✅ Chosen action: {action}")
                        return action
                    else:
                        # Try to find closest matching action
                        action_lower = action.lower()
                        for available in available_actions:
                            if available.lower() == action_lower:
                                if verbose:
                                    print(f"\n✅ Chosen action (corrected): {available}")
                                return available

            # Fallback if no valid action found
            print(f"⚠️ Warning: Could not parse valid action from LLM response. Using fallback.")
            return available_actions[0]

        except Exception as e:
            print(f"❌ Error calling LLM API: {e}")
            # Fallback to first available action
            return available_actions[0]

    def update_experience(self, state: str, action: str, feedback: str, reward: float):
        """
        Store an experience for future in-context learning.
        """
        # Determine if action was successful based on reward
        success = reward > 0

        experience = GameExperience(
            state_description=state,
            action=action,
            feedback=feedback,
            reward=reward,
            success=success
        )

        self.experiences.append(experience)

        # Keep only recent experiences to manage context length
        if len(self.experiences) > self.max_experiences * 2:
            # Keep a mix of successful and failed experiences
            successful = [e for e in self.experiences if e.success]
            failed = [e for e in self.experiences if not e.success]

            # Keep recent ones and some diverse older ones
            self.experiences = (
                successful[-self.max_experiences:] + 
                failed[-self.max_experiences//2:]
            )[-self.max_experiences:]

    def play_episode(self, game: TreasureHuntGame, verbose: bool = True) -> Tuple[float, int, bool]:
        """
        Play one episode of the game.
        """
        game.reset()
        total_reward = 0
        steps = 0
        trajectory = []

        if verbose:
            print("\n" + "🎮"*30)
            print("STARTING NEW GAME EPISODE")
            print("🎮"*30)

        while not game.game_over:
            if verbose:
                print(f"\n{'='*60}")
                print(f"STEP {steps + 1}")
                print(f"{'='*60}")

                # Show current game state
                print("\n📍 Current State:")
                state_lines = game.get_state_description().split('\n')
                for line in state_lines:
                    if line.strip():
                        print(f"  {line}")

            # Get state before action
            state_before = game.get_state_description()

            # Choose action using LLM
            action = self.choose_action(game, verbose=verbose)

            # Execute action
            feedback, reward, done = game.execute_action(action)

            # Store experience
            self.update_experience(state_before, action, feedback, reward)

            # Record trajectory
            trajectory.append({
                "step": steps + 1,
                "action": action,
                "reward": reward,
                "feedback": feedback
            })

            total_reward += reward
            steps += 1

            if verbose:
                print(f"\n🎯 Action Result:")
                print(f"  Feedback: {feedback}")
                if reward > 0:
                    print(f"  Reward: ✨ +{reward:.1f}")
                else:
                    print(f"  Reward: 📉 {reward:.1f}")
                print(f"  Total reward so far: {total_reward:.1f}")

                # Add a pause between steps for readability
                if not done:
                    print("\n" + "."*60)

        # Update statistics
        self.episode_rewards.append(total_reward)
        self.episode_lengths.append(steps)
        if game.victory:
            self.victories += 1
        self.total_episodes += 1

        if verbose:
            print("\n" + "🏁"*30)
            if game.victory:
                print("🎉 VICTORY! The LLM found the treasure!")
            else:
                print("💀 GAME OVER! Better luck next time.")
            print(f"  Final Score: {total_reward:.1f}")
            print(f"  Total Steps: {steps}")
            print(f"  API Calls Used: {self.api_calls}")
            print("🏁"*30)

        return total_reward, steps, game.victory

    def train(self, num_episodes: int = 20, verbose: bool = True, stochastic: bool = False) -> Dict[str, Any]:
        """
        'Train' the agent through in-context learning over multiple episodes.
        Note: Unlike traditional RL, there's no explicit training - just experience accumulation.

        Args:
            num_episodes: Number of episodes to play
            verbose: Whether to print details
            stochastic: Whether to use stochastic environment
        """
        game = TreasureHuntGame(stochastic=stochastic)

        print("\n" + "🚀"*30)
        print("LLM IN-CONTEXT LEARNING EXPERIMENT")
        print("🚀"*30)
        print(f"\n📝 Will play {num_episodes} episodes to learn the game")
        print("🧠 The LLM learns by accumulating experiences in context")
        print("⚡ Each decision shows the full reasoning process")

        for episode in range(num_episodes):
            print(f"\n\n{'🎯'*30}")
            print(f"EPISODE {episode + 1} of {num_episodes}")
            print(f"{'🎯'*30}")
            print(f"📚 Experiences accumulated so far: {len(self.experiences)}")

            # Show full process for first 3 episodes, then reduce verbosity
            show_full = verbose and (episode < 3 or episode == num_episodes - 1)

            if not show_full and verbose:
                print("\n(Reducing verbosity for middle episodes to save space...)")

            reward, steps, victory = self.play_episode(game, verbose=show_full)

            if not show_full:
                # Still show summary even when not fully verbose
                print(f"\n📊 Episode {episode + 1} Summary:")
                print(f"  Result: {'🎉 Victory!' if victory else '💀 Failed'}")
                print(f"  Total Reward: {reward:.2f}")
                print(f"  Steps Taken: {steps}")
                print(f"  Total API Calls So Far: {self.api_calls}")

            # Show learning progress
            if len(self.episode_rewards) >= 3:
                recent_victories = sum(1 for r in self.episode_rewards[-3:] if r > 50)
                recent_avg = sum(self.episode_rewards[-3:]) / 3
                print(f"\n📈 Recent Performance (last 3 episodes):")
                print(f"  Victories: {recent_victories}/3")
                print(f"  Average Reward: {recent_avg:.2f}")

            # Add delay to respect rate limits
            if episode < num_episodes - 1:
                print("\n⏳ Waiting 1 second for API rate limits...")
                time.sleep(1)

        return {
            "total_episodes": self.total_episodes,
            "total_victories": self.victories,
            "victory_rate": self.victories / self.total_episodes if self.total_episodes > 0 else 0,
            "total_api_calls": self.api_calls,
            "total_tokens": self.total_tokens,
            "experiences_collected": len(self.experiences),
            "episode_rewards": self.episode_rewards,
            "episode_lengths": self.episode_lengths
        }

    def evaluate(self, num_episodes: int = 10, verbose: bool = False, stochastic: bool = False) -> Dict[str, Any]:
        """
        Evaluate the agent's performance using accumulated experiences.

        Args:
            num_episodes: Number of episodes to evaluate
            verbose: Whether to print details
            stochastic: Whether to use stochastic environment
        """
        game = TreasureHuntGame(stochastic=stochastic)
        eval_rewards = []
        eval_lengths = []
        eval_victories = 0

        for episode in range(num_episodes):
            reward, steps, victory = self.play_episode(game, verbose=verbose)

            eval_rewards.append(reward)
            eval_lengths.append(steps)
            if victory:
                eval_victories += 1

            if verbose:
                print(f"Episode {episode + 1}: Reward={reward:.2f}, Steps={steps}, Victory={victory}")

        return {
            "num_episodes": num_episodes,
            "victories": eval_victories,
            "victory_rate": eval_victories / num_episodes if num_episodes else 0.0,
            "avg_reward": sum(eval_rewards) / len(eval_rewards) if eval_rewards else 0.0,
            "avg_length": sum(eval_lengths) / len(eval_lengths) if eval_lengths else 0.0,
            "total_api_calls": self.api_calls,
            "experiences_used": len(self.experiences)
        }

    def save_experiences(self, filepath: str):
        """Save experiences to file for analysis."""
        data = {
            "experiences": [asdict(exp) for exp in self.experiences],
            "statistics": {
                "total_episodes": self.total_episodes,
                "victories": self.victories,
                "api_calls": self.api_calls,
                "total_tokens": self.total_tokens
            }
        }

        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)

    def load_experiences(self, filepath: str):
        """Load experiences from file."""
        with open(filepath, 'r') as f:
            data = json.load(f)

        self.experiences = [
            GameExperience(**exp) for exp in data["experiences"]
        ]

        stats = data.get("statistics", {})
        self.total_episodes = stats.get("total_episodes", 0)
        self.victories = stats.get("victories", 0)
        self.api_calls = stats.get("api_calls", 0)
        self.total_tokens = stats.get("total_tokens", 0)

quick_demo.py

#!/usr/bin/env python3
"""
Quick demo showing the LLM learning process in detail.
This script runs a simplified experiment to demonstrate how LLMs learn from experience.
"""

import os
import sys
from pathlib import Path

from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))

from game_environment import TreasureHuntGame
from llm_agent import LLMAgent


def show_game_solution():
    """Show the optimal solution to the game."""
    print("\n" + "="*70)
    print("GAME SOLUTION (for reference)")
    print("="*70)

    game = TreasureHuntGame()
    print(game.get_hidden_rules())

    print("\n📝 Optimal solution path:")
    print("1. Take rusty sword (in entrance)")
    print("2. Go east to storage")
    print("3. Take red key")
    print("4. Take magic crystal")
    print("5. Try crafting → creates silver sword")
    print("6. Go west to entrance")
    print("7. Go north to hallway (uses red key automatically)")
    print("8. Go north to guard room")
    print("9. Attack with silver sword → defeats strong guard")
    print("10. Go east to treasure room")
    print("11. Take dragon's treasure → Victory!")
    print("\n✨ Total moves: ~11-12 (optimal)")


def run_llm_demo():
    """Run a simplified LLM demo with just a few episodes."""
    print("\n" + "🤖"*35)
    print("LLM IN-CONTEXT LEARNING DEMO")
    print("🤖"*35)

    # Check API key
    api_key = os.getenv("MOONSHOT_API_KEY")
    if not api_key and not os.getenv("OPENROUTER_API_KEY"):
        print("\n❌ Error: MOONSHOT_API_KEY not set.")
        print("Please set your Kimi API key:")
        print("  export MOONSHOT_API_KEY='your-key-here'")
        print("\nGet your key at: https://platform.moonshot.cn/")
        print("Or set OPENROUTER_API_KEY as a universal fallback.")
        return

    print("\n✅ API key found!")
    print("🧠 Initializing Kimi K3 LLM agent...")

    # Initialize agent
    agent = LLMAgent(
        api_key=api_key,
        model=os.getenv("MOONSHOT_MODEL", "kimi-k3"),
        temperature=0.7,
        max_experiences=30
    )

    print("\n📚 The LLM will play 3 episodes to learn the game")
    print("👀 Watch how it reasons and learns from each experience!\n")

    # Play 3 episodes
    game = TreasureHuntGame()

    for episode in range(3):
        print("\n" + "🎮"*35)
        print(f"EPISODE {episode + 1} of 3")
        print("🎮"*35)

        # Show what the LLM has learned so far
        if agent.experiences:
            print(f"\n📊 Experience Memory: {len(agent.experiences)} interactions stored")

            # Show some key learnings
            successful = [e for e in agent.experiences if e.success]
            if successful:
                print("✅ Successful patterns discovered:")
                for exp in successful[-3:]:
                    print(f"   • {exp.action} → reward: {exp.reward:.1f}")

            failed = [e for e in agent.experiences if not e.success]
            if failed and len(failed) > 5:
                print("❌ Mistakes to avoid:")
                for exp in failed[-2:]:
                    print(f"   • {exp.action} → reward: {exp.reward:.1f}")

        # Play episode
        reward, steps, victory = agent.play_episode(game, verbose=True)

        print(f"\n📈 Episode {episode + 1} Performance:")
        print(f"  • Result: {'🎉 Victory!' if victory else '💀 Failed'}")
        print(f"  • Total Reward: {reward:.2f}")
        print(f"  • Steps Taken: {steps}")
        print(f"  • Experiences Collected: {len(agent.experiences)}")

        if victory:
            print("\n🎊 The LLM learned to solve the game!")
            print(f"   It took {episode + 1} episodes to learn")
            print(f"   Total API calls used: {agent.api_calls}")
            break

        if episode < 2:
            print("\n⏳ Waiting 2 seconds before next episode...")
            import time
            time.sleep(2)

    # Summary
    print("\n" + "="*70)
    print("DEMO SUMMARY")
    print("="*70)
    print(f"📊 Total episodes played: {episode + 1}")
    print(f"🧠 Total experiences collected: {len(agent.experiences)}")
    print(f"🎯 Victories: {agent.victories}")
    print(f"📡 API calls made: {agent.api_calls}")

    if agent.victories > 0:
        print("\n✨ Key Insight:")
        print("The LLM learned to solve the game by reasoning about patterns")
        print("in just a few episodes, without any parameter updates!")
        print("Traditional RL would need thousands of episodes for the same result.")
    else:
        print("\n💡 Note: The LLM is still learning. Run more episodes to see it succeed!")


def main():
    """Main entry point."""
    print("\n" + "🎯"*35)
    print("LEARNING FROM EXPERIENCE: LLM DEMO")
    print("Replicating insights from 'The Second Half'")
    print("🎯"*35)

    # Show solution first
    show_game_solution()

    # Ask user if they want to continue
    response = input("\n▶️ Ready to see how an LLM learns this game? (y/n): ").strip().lower()

    if response == 'y':
        run_llm_demo()
    else:
        print("\n👋 Okay, goodbye!")


if __name__ == "__main__":
    main()

rl_agent.py

"""
Traditional Reinforcement Learning Agent using Q-learning.
This demonstrates the classical RL approach that requires extensive training.
"""

import numpy as np
import pickle
from collections import defaultdict
from typing import Dict, List, Tuple, Any
import random
from game_environment import TreasureHuntGame


class QLearningAgent:
    """
    Q-learning agent for the treasure hunt game.
    Uses tabular Q-learning with state-action pairs.
    """

    def __init__(self, 
                 learning_rate: float = 0.2,
                 discount_factor: float = 0.99,
                 epsilon: float = 1.0,
                 epsilon_decay: float = 0.9995,
                 epsilon_min: float = 0.1):
        """
        Initialize Q-learning agent.

        Args:
            learning_rate: Alpha parameter for Q-value updates
            discount_factor: Gamma parameter for future rewards
            epsilon: Initial exploration rate
            epsilon_decay: Rate at which epsilon decreases
            epsilon_min: Minimum exploration rate
        """
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min

        # Q-table: state_hash -> action -> Q-value
        self.q_table = defaultdict(lambda: defaultdict(float))

        # Statistics
        self.episode_rewards = []
        self.episode_lengths = []
        self.episode_victories = []  # Per-episode victory flag (1/0), for learning curves
        self.victories = 0
        self.total_episodes = 0
        self.learning_curve = []  # Snapshots recorded at checkpoints during train()

    def _get_state_hash(self, game: TreasureHuntGame) -> str:
        """
        Create a hashable representation of the game state.
        This is crucial for tabular Q-learning.
        """
        # Include relevant state information
        state_parts = [
            game.current_room.name,
            tuple(sorted([item.name for item in game.inventory])),
            tuple(sorted([item.name for item in game.current_room.items])),
            tuple(sorted(game.current_room.locked_exits.items())),
            game.current_room.has_guard and not game.current_room.guard_defeated
        ]

        return str(state_parts)

    def choose_action(self, game: TreasureHuntGame, training: bool = True) -> str:
        """
        Choose an action using epsilon-greedy strategy.
        """
        available_actions = game.get_available_actions()

        if not available_actions:
            return "look around"

        # Exploration vs exploitation
        if training and random.random() < self.epsilon:
            # Explore: choose random action
            return random.choice(available_actions)
        else:
            # Exploit: choose best action based on Q-values
            state_hash = self._get_state_hash(game)

            # Get Q-values for all available actions
            action_values = {
                action: self.q_table[state_hash][action] 
                for action in available_actions
            }

            # If all Q-values are 0 (unexplored), choose randomly
            if all(v == 0 for v in action_values.values()):
                return random.choice(available_actions)

            # Choose action with highest Q-value
            return max(action_values, key=action_values.get)

    def update_q_value(self, state: str, action: str, reward: float, 
                       next_state: str, next_actions: List[str], done: bool):
        """
        Update Q-value using the Q-learning update rule.
        Q(s,a) <- Q(s,a) + α[r + γ max Q(s',a') - Q(s,a)]
        """
        current_q = self.q_table[state][action]

        if done:
            # Terminal state
            target = reward
        else:
            # Get maximum Q-value for next state
            if next_actions:
                max_next_q = max(
                    self.q_table[next_state][a] for a in next_actions
                )
            else:
                max_next_q = 0

            target = reward + self.discount_factor * max_next_q

        # Update Q-value
        self.q_table[state][action] = (
            current_q + self.learning_rate * (target - current_q)
        )

    def train_episode(self, game: TreasureHuntGame) -> Tuple[float, int, bool]:
        """
        Train the agent for one episode.

        Returns:
            Total reward, number of steps, victory status
        """
        game.reset()
        total_reward = 0
        steps = 0

        while not game.game_over:
            # Get current state
            state_hash = self._get_state_hash(game)

            # Choose action
            action = self.choose_action(game, training=True)

            # Execute action
            feedback, reward, done = game.execute_action(action)

            # Get next state
            next_state_hash = self._get_state_hash(game)
            next_actions = game.get_available_actions() if not done else []

            # Update Q-value
            self.update_q_value(
                state_hash, action, reward, 
                next_state_hash, next_actions, done
            )

            total_reward += reward
            steps += 1

        # Decay epsilon
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

        # Update statistics
        self.episode_rewards.append(total_reward)
        self.episode_lengths.append(steps)
        self.episode_victories.append(1 if game.victory else 0)
        if game.victory:
            self.victories += 1
        self.total_episodes += 1

        return total_reward, steps, game.victory

    def train(self, num_episodes: int = 1000, verbose: bool = True,
              stochastic: bool = False, checkpoint_interval: int = 0) -> Dict[str, Any]:
        """
        Train the agent for multiple episodes.

        Args:
            num_episodes: Number of episodes to train
            verbose: Whether to print progress
            stochastic: Whether to use stochastic environment
            checkpoint_interval: If > 0, record a learning-curve snapshot
                (episode, windowed victory rate, Q-table size, epsilon) every
                this many episodes. Snapshots are stored in self.learning_curve.
        """
        game = TreasureHuntGame(stochastic=stochastic)

        # Adjust hyperparameters for stochastic environment
        if stochastic:
            # Slightly slower epsilon decay for stochastic environments
            original_decay = self.epsilon_decay
            self.epsilon_decay = min(0.9999, self.epsilon_decay * 1.001)
            if verbose:
                print(f"Adjusted epsilon_decay from {original_decay:.4f} to {self.epsilon_decay:.4f} for stochastic environment\n")

        window = checkpoint_interval if checkpoint_interval and checkpoint_interval > 0 else 1000

        for episode in range(num_episodes):
            reward, steps, victory = self.train_episode(game)

            # Record a learning-curve snapshot at each checkpoint
            if checkpoint_interval and checkpoint_interval > 0 and (episode + 1) % checkpoint_interval == 0:
                recent = self.episode_victories[-window:]
                self.learning_curve.append({
                    "episode": episode + 1,
                    "victory_rate": sum(recent) / len(recent),
                    "q_table_size": len(self.q_table),
                    "epsilon": self.epsilon,
                })

            if verbose and (episode + 1) % 100 == 0:
                recent_rewards = self.episode_rewards[-100:]
                recent_victories = sum(
                    1 for r in recent_rewards if r > 50  # Approximate victory
                )
                avg_reward = np.mean(recent_rewards)

                print(f"Episode {episode + 1}/{num_episodes}")
                print(f"  Avg Reward (last 100): {avg_reward:.2f}")
                print(f"  Victories (last 100): {recent_victories}")
                print(f"  Epsilon: {self.epsilon:.3f}")
                print(f"  Q-table size: {len(self.q_table)}")
                print()

        return {
            "total_episodes": self.total_episodes,
            "total_victories": self.victories,
            "victory_rate": self.victories / self.total_episodes if self.total_episodes else 0.0,
            "final_epsilon": self.epsilon,
            "q_table_size": len(self.q_table),
            "episode_rewards": self.episode_rewards,
            "episode_lengths": self.episode_lengths,
            "learning_curve": self.learning_curve,
        }

    def evaluate(self, num_episodes: int = 100, verbose: bool = False, stochastic: bool = False) -> Dict[str, Any]:
        """
        Evaluate the trained agent without learning.

        Args:
            num_episodes: Number of episodes to evaluate
            verbose: Whether to print details
            stochastic: Whether to use stochastic environment
        """
        game = TreasureHuntGame(stochastic=stochastic)
        eval_rewards = []
        eval_lengths = []
        eval_victories = 0

        # Store original epsilon and set to 0 for evaluation
        original_epsilon = self.epsilon
        self.epsilon = 0

        for episode in range(num_episodes):
            game.reset()
            total_reward = 0
            steps = 0

            while not game.game_over:
                action = self.choose_action(game, training=False)
                feedback, reward, done = game.execute_action(action)
                total_reward += reward
                steps += 1

                if verbose and episode == 0:  # Show first evaluation episode
                    print(f"Step {steps}: {action}")
                    print(f"Feedback: {feedback}")
                    print()

            eval_rewards.append(total_reward)
            eval_lengths.append(steps)
            if game.victory:
                eval_victories += 1

        # Restore epsilon
        self.epsilon = original_epsilon

        return {
            "num_episodes": num_episodes,
            "victories": eval_victories,
            "victory_rate": eval_victories / num_episodes if num_episodes else 0.0,
            "avg_reward": np.mean(eval_rewards),
            "std_reward": np.std(eval_rewards),
            "avg_length": np.mean(eval_lengths),
            "std_length": np.std(eval_lengths)
        }

    def save(self, filepath: str):
        """Save the Q-table and parameters."""
        data = {
            "q_table": dict(self.q_table),
            "epsilon": self.epsilon,
            "learning_rate": self.learning_rate,
            "discount_factor": self.discount_factor,
            "statistics": {
                "total_episodes": self.total_episodes,
                "victories": self.victories,
                "episode_rewards": self.episode_rewards,
                "episode_lengths": self.episode_lengths
            }
        }

        with open(filepath, 'wb') as f:
            pickle.dump(data, f)

    def load(self, filepath: str):
        """Load a saved Q-table and parameters."""
        with open(filepath, 'rb') as f:
            data = pickle.load(f)

        self.q_table = defaultdict(lambda: defaultdict(float))
        for state, actions in data["q_table"].items():
            for action, value in actions.items():
                self.q_table[state][action] = value

        self.epsilon = data["epsilon"]
        self.learning_rate = data["learning_rate"]
        self.discount_factor = data["discount_factor"]

        stats = data.get("statistics", {})
        self.total_episodes = stats.get("total_episodes", 0)
        self.victories = stats.get("victories", 0)
        self.episode_rewards = stats.get("episode_rewards", [])
        self.episode_lengths = stats.get("episode_lengths", [])


class DQNAgent:
    """
    Deep Q-Network agent for comparison.
    Uses neural network function approximation instead of tabular Q-learning.
    """

    def __init__(self, 
                 state_dim: int = 128,
                 hidden_dim: int = 256,
                 learning_rate: float = 0.001,
                 discount_factor: float = 0.95,
                 epsilon: float = 1.0,
                 epsilon_decay: float = 0.995,
                 epsilon_min: float = 0.01,
                 batch_size: int = 32,
                 memory_size: int = 10000):
        """
        Initialize DQN agent with neural network.
        Note: Simplified implementation for demonstration.
        """
        self.state_dim = state_dim
        self.hidden_dim = hidden_dim
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min
        self.batch_size = batch_size

        # Experience replay buffer
        self.memory = []
        self.memory_size = memory_size

        # Statistics
        self.episode_rewards = []
        self.episode_lengths = []
        self.victories = 0
        self.total_episodes = 0

        # Note: For full implementation, we would use PyTorch or TensorFlow
        # This is a simplified placeholder
        print("Note: DQN implementation requires neural network library.")
        print("Using simplified random policy for demonstration.")

    def choose_action(self, game: TreasureHuntGame, training: bool = True) -> str:
        """Choose action (simplified for demonstration)."""
        available_actions = game.get_available_actions()
        if not available_actions:
            return "look around"

        # Simplified: just use epsilon-greedy with random selection
        if training and random.random() < self.epsilon:
            return random.choice(available_actions)
        else:
            # In full implementation, this would use neural network
            return random.choice(available_actions)

    def train_episode(self, game: TreasureHuntGame) -> Tuple[float, int, bool]:
        """Train for one episode (simplified)."""
        game.reset()
        total_reward = 0
        steps = 0

        while not game.game_over:
            action = self.choose_action(game, training=True)
            feedback, reward, done = game.execute_action(action)
            total_reward += reward
            steps += 1

        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

        self.episode_rewards.append(total_reward)
        self.episode_lengths.append(steps)
        if game.victory:
            self.victories += 1
        self.total_episodes += 1

        return total_reward, steps, game.victory

test_basic.py

#!/usr/bin/env python3
"""
Basic test to verify all components work correctly.
"""

import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))


def test_game_environment():
    """Test that the game environment works."""
    print("Testing game environment...")
    from game_environment import TreasureHuntGame

    game = TreasureHuntGame(seed=42)

    # Test initial state
    state = game.get_state_description()
    assert "entrance" in state.lower()
    print("  ✓ Game initialization works")

    # Test actions
    actions = game.get_available_actions()
    assert len(actions) > 0
    print("  ✓ Actions generation works")

    # Test action execution
    feedback, reward, done = game.execute_action("look around")
    assert isinstance(feedback, str)
    assert isinstance(reward, float)
    assert isinstance(done, bool)
    print("  ✓ Action execution works")

    # Test reset
    game.reset()
    assert game.moves == 0
    print("  ✓ Game reset works")

    print("✅ Game environment tests passed!\n")


def test_rl_agent():
    """Test that the RL agent works."""
    print("Testing RL agent...")
    from game_environment import TreasureHuntGame
    from rl_agent import QLearningAgent

    game = TreasureHuntGame(seed=42)
    agent = QLearningAgent()

    # Test action selection
    action = agent.choose_action(game, training=True)
    assert isinstance(action, str)
    print("  ✓ Action selection works")

    # Test Q-value update
    state = agent._get_state_hash(game)
    feedback, reward, done = game.execute_action(action)
    next_state = agent._get_state_hash(game)
    next_actions = game.get_available_actions()

    agent.update_q_value(state, action, reward, next_state, next_actions, done)
    print("  ✓ Q-value update works")

    # Test training (just 10 episodes for speed)
    results = agent.train(num_episodes=10, verbose=False)
    assert "total_episodes" in results
    print("  ✓ Training works")

    print("✅ RL agent tests passed!\n")


def test_llm_agent():
    """Test that the LLM agent works (without API calls)."""
    print("Testing LLM agent structure...")
    from game_environment import TreasureHuntGame
    from llm_agent import LLMAgent, GameExperience

    # Test experience storage
    exp = GameExperience(
        state_description="test state",
        action="test action",
        feedback="test feedback",
        reward=1.0,
        success=True
    )
    assert exp.action == "test action"
    print("  ✓ Experience dataclass works")

    # Test context building (without API)
    try:
        # This will fail without API key, but we can test the structure
        agent = LLMAgent(api_key="dummy-key-for-testing")

        game = TreasureHuntGame()
        state = game.get_state_description()
        actions = game.get_available_actions()

        context = agent._build_context(state, actions)
        assert "treasure hunt" in context.lower()
        print("  ✓ Context building works")

        # Test experience update
        agent.update_experience(state, "test action", "test feedback", 1.0)
        assert len(agent.experiences) == 1
        print("  ✓ Experience storage works")

    except ValueError as e:
        if "MOONSHOT_API_KEY" in str(e):
            print("  ⚠ LLM agent requires API key for full testing")
        else:
            raise

    print("✅ LLM agent structure tests passed!\n")


def test_experiment_runner():
    """Test that the experiment runner works."""
    print("Testing experiment runner...")
    from experiment import ExperimentRunner

    runner = ExperimentRunner(results_dir="test_results")
    assert runner.results_dir.exists()
    print("  ✓ Experiment runner initialization works")

    # Clean up test directory
    import shutil
    if runner.results_dir.exists():
        shutil.rmtree(runner.results_dir)

    print("✅ Experiment runner tests passed!\n")


def main():
    """Run all tests."""
    print("\n" + "="*60)
    print("RUNNING BASIC TESTS")
    print("="*60 + "\n")

    try:
        test_game_environment()
        test_rl_agent()
        test_llm_agent()
        test_experiment_runner()

        print("="*60)
        print("ALL TESTS PASSED! ✅")
        print("="*60)
        print("\nThe experiment is ready to run.")
        print("To run the full experiment: python experiment.py")
        print("To play interactively: python demo.py")

    except Exception as e:
        print(f"\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())

test_rl_learning.py

#!/usr/bin/env python3
"""
Quick test to verify Q-learning agent can learn the simplified game.
"""

import sys
import argparse
from pathlib import Path
import numpy as np

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))

from game_environment import TreasureHuntGame
from rl_agent import QLearningAgent


def test_rl_learning(stochastic=False, episodes=None):
    """Test that Q-learning can learn the game.

    Args:
        stochastic: If True, use stochastic environment
        episodes: List of episode counts to test (default: various counts)
    """
    env_type = "STOCHASTIC" if stochastic else "DETERMINISTIC"
    print(f"Testing Q-Learning on simplified game ({env_type} environment)...")
    print("="*50)

    # Show game rules
    game = TreasureHuntGame(stochastic=stochastic)
    print(game.get_hidden_rules())
    if stochastic:
        print("\n⚠️  Stochastic Mode Active:")
        print("  - Random reward variations")
        print("  - 3% chance of action failure")
        print("  - 10% critical hit / 5% miss chance in combat")
        print("  - 10% crafting failure chance")
    print("\n" + "="*50)

    # Initialize agent
    agent = QLearningAgent(
        learning_rate=0.2,
        discount_factor=0.99,
        epsilon=1.0,
        epsilon_decay=0.9997,  # Slower decay for exploration
        epsilon_min=0.1
    )

    # Train for different episode counts
    if episodes:
        episode_counts = episodes
    else:
        episode_counts = [100, 500, 1000, 2000, 5000, 10000]

    for num_episodes in episode_counts:
        print(f"\nTraining for {num_episodes} episodes...")

        # Reset agent
        agent = QLearningAgent(
            learning_rate=0.2,
            discount_factor=0.99,
            epsilon=1.0,
            epsilon_decay=0.9997,
            epsilon_min=0.1
        )

        # Train
        game = TreasureHuntGame(stochastic=stochastic)
        victories = 0
        recent_rewards = []

        for episode in range(num_episodes):
            game.reset()
            total_reward = 0

            while not game.game_over:
                state_hash = agent._get_state_hash(game)
                action = agent.choose_action(game, training=True)
                feedback, reward, done = game.execute_action(action)

                next_state_hash = agent._get_state_hash(game)
                next_actions = game.get_available_actions() if not done else []

                agent.update_q_value(
                    state_hash, action, reward, 
                    next_state_hash, next_actions, done
                )

                total_reward += reward

            # Decay epsilon
            agent.epsilon = max(agent.epsilon_min, agent.epsilon * agent.epsilon_decay)

            recent_rewards.append(total_reward)
            if game.victory:
                victories += 1

            # Print progress
            if (episode + 1) % (num_episodes // 10) == 0:
                recent_wins = sum(1 for r in recent_rewards[-100:] if r > 50)
                avg_reward = np.mean(recent_rewards[-100:]) if recent_rewards else 0
                print(f"  Episode {episode+1}: Recent wins={recent_wins}/100, "
                      f"Avg reward={avg_reward:.1f}, Epsilon={agent.epsilon:.3f}")

        # Evaluate
        print(f"\nEvaluating after {num_episodes} episodes...")
        eval_victories = 0
        eval_rewards = []

        for _ in range(100):
            game.reset()
            total_reward = 0

            # Set epsilon to 0 for evaluation
            old_epsilon = agent.epsilon
            agent.epsilon = 0

            while not game.game_over:
                action = agent.choose_action(game, training=False)
                feedback, reward, done = game.execute_action(action)
                total_reward += reward

            agent.epsilon = old_epsilon

            eval_rewards.append(total_reward)
            if game.victory:
                eval_victories += 1

        print(f"  Evaluation: {eval_victories}/100 victories")
        print(f"  Average reward: {np.mean(eval_rewards):.2f}")
        print(f"  Q-table size: {len(agent.q_table)} states")

        # Show a sample successful trajectory if we have victories
        if eval_victories > 0:
            print("\n  Sample successful trajectory:")
            game.reset()
            agent.epsilon = 0
            steps = []

            while not game.game_over:
                action = agent.choose_action(game, training=False)
                steps.append(f"    {len(steps)+1}. {action}")
                feedback, reward, done = game.execute_action(action)
                if game.victory:
                    steps.append(f"    → Victory! Total moves: {game.moves}")
                    break

            if len(steps) <= 20:  # Only show if reasonable length
                print("\n".join(steps[:15]))  # Show first 15 steps


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test Q-learning agent on the treasure hunt game")
    parser.add_argument(
        '--stochastic', 
        action='store_true',
        help='Use stochastic environment (adds randomness to rewards and actions)'
    )
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='Use deterministic environment (default)'
    )
    parser.add_argument(
        '--episodes',
        type=int,
        nargs='+',
        help='Episode counts to test (e.g., --episodes 1000 5000 10000)'
    )

    args = parser.parse_args()

    # Handle environment mode
    if args.deterministic and args.stochastic:
        print("Error: Cannot specify both --deterministic and --stochastic")
        sys.exit(1)

    stochastic = args.stochastic  # Default is False (deterministic)

    test_rl_learning(stochastic=stochastic, episodes=args.episodes)

test_zero_episodes.py

#!/usr/bin/env python3
"""Regression tests for zero-episode division guards.

Bug: train()/evaluate() divided victory counts by episode counts, so
num_episodes=0 (accepted by experiment.py's argparse) crashed with
ZeroDivisionError. Fixed by guarding the divisions and rejecting
episode counts < 1 in experiment.py's front door.
"""

import sys

import experiment
from llm_agent import LLMAgent
from rl_agent import QLearningAgent


def test_rl_train_zero_episodes_no_zero_division():
    result = QLearningAgent().train(num_episodes=0, verbose=False)
    assert result["total_episodes"] == 0
    assert result["victory_rate"] == 0.0


def test_rl_evaluate_zero_episodes_no_zero_division():
    result = QLearningAgent().evaluate(num_episodes=0)
    assert result["num_episodes"] == 0
    assert result["victory_rate"] == 0.0


def test_llm_evaluate_zero_episodes_no_zero_division():
    # Dummy key: constructing the client makes no network calls, and
    # evaluate(num_episodes=0) never reaches the API.
    agent = LLMAgent(api_key="dummy-key")
    result = agent.evaluate(num_episodes=0)
    assert result["victory_rate"] == 0.0
    assert result["avg_reward"] == 0.0
    assert result["avg_length"] == 0.0


def test_experiment_rejects_zero_episodes(monkeypatch, capsys):
    monkeypatch.setattr(sys, "argv", ["experiment.py", "--mode", "qlearning",
                                      "--rl-episodes", "0"])
    experiment.main()  # must print an error and return before running
    assert "must all be >= 1" in capsys.readouterr().out