Upload folder using huggingface_hub

README.md

@@ -12,8 +12,83 @@ tags:
   - reinforcement-learning
   - planning
   - ai-education
+  - model-based-rl
+  - muzero
+  - dreamer
 ---
 
-# World Model Demo
+# 🧠 World Model Demo
 
-Interactive visualization of model-based reinforcement learning concepts.
+**An interactive visualization of model-based reinforcement learning concepts**
+
+## What is a World Model?
+
+A **world model** is an internal representation that an AI agent uses to *simulate* the environment without actually interacting with it. Think of it as the agent's "imagination" - it can mentally rehearse actions and predict their outcomes before committing to them in the real world.
+
+### The Key Insight
+
+Instead of learning through pure trial-and-error (which is slow and potentially dangerous), an agent with a world model can:
+
+1. **Imagine** possible futures by simulating "what if I do X?"
+2. **Evaluate** which imagined future looks best
+3. **Plan** a sequence of actions to reach that future
+4. **Act** with confidence, having already "seen" the outcome
+
+## How This Differs from Language Models
+
+| Aspect | Language Model (GPT, Claude) | World Model (MuZero, Dreamer) |
+|--------|------------------------------|-------------------------------|
+| **Primary function** | Predict next token in a sequence | Predict next *state* given an action |
+| **Training signal** | Text prediction loss | Reward from environment |
+| **"Imagination"** | Generates plausible text continuations | Simulates future environment states |
+| **Planning** | Implicit (via chain-of-thought) | Explicit (via tree search or rollouts) |
+| **Grounding** | Statistical patterns in text | Causal dynamics of an environment |
+
+### A Concrete Example
+
+**Language Model**: "If I push a ball off a table, it will..." → generates plausible text based on patterns
+
+**World Model**: Given state (ball on table) + action (push) → predicts new state (ball falling, trajectory, landing position) with enough fidelity to *plan* around it
+
+## What You're Seeing in This Demo
+
+This visualization shows a simplified world model operating on a grid navigation task:
+
+### The Four Phases
+
+1. **🔍 Observe**: The agent perceives the current grid state (its position, goal location, obstacles)
+
+2. **💭 Imagine**: The world model predicts what would happen for each possible action (up/down/left/right). You see this as the "mental simulation" exploring future states.
+
+3. **🌳 Plan**: Using tree search (similar to how chess engines work), the agent evaluates sequences of actions by imagining multiple steps ahead. Better paths to the goal get higher scores.
+
+4. **⚡ Act**: The agent executes the best action found during planning, then the cycle repeats.
+
+### Why This Matters for AI Safety
+
+World models are crucial for AI safety research because:
+
+- **Predictability**: Agents that plan can be analyzed - we can inspect what futures they're considering
+- **Corrigibility**: Planning agents can incorporate "don't do irreversible things" into their search
+- **Interpretability**: The world model's predictions can be examined for accuracy and bias
+- **Scalable oversight**: Humans can audit the agent's "reasoning" by inspecting its simulated futures
+
+## Real-World Architectures
+
+This demo is inspired by:
+
+- **MuZero** (DeepMind): Learned world models that mastered Go, chess, and Atari without knowing the rules
+- **Dreamer** (Hafner et al.): World models for continuous control from pixels
+- **IRIS** (Micheli et al.): Transformer-based world models for Atari
+- **Genie** (DeepMind): Generative world models from video
+
+## Try It Yourself
+
+1. Click **"Run World Model"** to watch the full planning cycle
+2. Use **Step Mode** to see each phase individually  
+3. Adjust grid size and obstacles to see how planning adapts
+4. Watch the **Imagined Futures** panel to see the agent's "thoughts"
+
+---
+
+*Created by [Anthony Maio](https://huggingface.co/anthonym21) as an educational resource for AI safety research*

app.py

@@ -99,21 +99,21 @@ class WorldModel:
 # Visualization
 # ============================================================================
 
-def render_grid_html(state, prediction=None, phase="observe"):
-    """Render the grid as an HTML table"""
-    size = state['size']
+def render_grid_html(state, phase="observe", prediction=None):
+    """Render the grid as HTML with phase-appropriate styling"""
     agent = state['agent']
     goal = state['goal']
-    obstacles = set(map(tuple, state['obstacles']))
+    obstacles = set(tuple(o) if isinstance(o, list) else o for o in state['obstacles'])
+    size = state['size']
     
-    colors = {
-        'observe': '#3b82f6',
-        'predict': '#8b5cf6', 
-        'plan': '#f59e0b',
-        'act': '#22c55e',
-        'learn': '#ec4899'
+    phase_colors = {
+        'observe': '#3b82f6',  # blue
+        'predict': '#f59e0b',  # amber
+        'plan': '#8b5cf6',     # purple
+        'act': '#10b981',      # green
+        'learn': '#ec4899'     # pink
     }
-    phase_color = colors.get(phase, '#6b7280')
+    phase_color = phase_colors.get(phase, '#6b7280')
     
     html = f'''
     <div style="text-align: center; font-family: system-ui, sans-serif;">
@@ -203,8 +203,9 @@ def do_action(action):
 
 def reset_env():
     global current_state, current_phase, world, model
+    world = GridWorld()  # Create fresh world
+    model = WorldModel()  # Create fresh model
     current_state = world.reset()
-    model = WorldModel()
     current_phase = "observe"
     html, stats = get_display()
     return html, stats, "Environment reset!"
@@ -214,14 +215,8 @@ with gr.Blocks(title="World Model Demo", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🧠 World Model Demo
     
-    Interactive demonstration of how AI agents build internal models of the world.
-    
-    **The Learning Cycle:**
-    1. **Observe** - Agent perceives current state
-    2. **Predict** - World model predicts action outcomes  
-    3. **Plan** - Agent evaluates possible futures
-    4. **Act** - Execute chosen action
-    5. **Learn** - Update model from observed outcome
+    **What is this?** An interactive demonstration of how AI agents can build internal "mental models" 
+    of the world to plan and reason, rather than just reacting to inputs.
     """)
     
     with gr.Row():
@@ -233,9 +228,9 @@ with gr.Blocks(title="World Model Demo", theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=1):
             gr.Markdown("### Controls")
             with gr.Row():
-                gr.Button("").click(lambda: None)
+                gr.Button("", visible=False, min_width=1)
                 up_btn = gr.Button("⬆️ Up")
-                gr.Button("").click(lambda: None)
+                gr.Button("", visible=False, min_width=1)
             with gr.Row():
                 left_btn = gr.Button("⬅️ Left")
                 down_btn = gr.Button("⬇️ Down")
@@ -244,16 +239,62 @@ with gr.Blocks(title="World Model Demo", theme=gr.themes.Soft()) as demo:
             reset_btn = gr.Button("🔄 Reset", variant="secondary")
             
             gr.Markdown("""
-            ### About World Models
-            
-            World models are internal representations that AI agents use to:
-            - Simulate possible futures
-            - Plan without trial-and-error
-            - Learn efficiently from experience
-            
-            Used in: MuZero, Dreamer, PlaNet
+            ---
+            **The Learning Cycle:**
+            1. 🔍 **Observe** - Perceive state
+            2. 💭 **Predict** - Imagine outcomes  
+            3. ⚡ **Act** - Execute action
+            4. 📚 **Learn** - Update model
             """)
     
+    # Educational content in collapsible sections
+    with gr.Accordion("📖 What is a World Model?", open=False):
+        gr.Markdown("""
+        A **world model** is an internal representation that an AI agent uses to *simulate* the 
+        environment without actually interacting with it. Think of it as the agent's "imagination."
+        
+        **Instead of pure trial-and-error, an agent with a world model can:**
+        - 🎯 **Imagine** possible futures ("what if I do X?")
+        - ⚖️ **Evaluate** which imagined future looks best
+        - 🗺️ **Plan** a sequence of actions to reach that future
+        - ✅ **Act** with confidence, having already "seen" the outcome
+        
+        **Real examples:** MuZero (mastered Go/Chess without knowing rules), Dreamer (robot control), 
+        IRIS (Atari from pixels)
+        """)
+    
+    with gr.Accordion("🤔 How is this different from ChatGPT/Claude?", open=False):
+        gr.Markdown("""
+        | Aspect | Language Model (GPT, Claude) | World Model (This Demo) |
+        |--------|------------------------------|-------------------------|
+        | **Predicts** | Next *word* in a sequence | Next *state* given an action |
+        | **Training** | Text prediction | Reward from environment |
+        | **"Thinking"** | Generates plausible text | Simulates physical outcomes |
+        | **Planning** | Implicit (chain-of-thought) | Explicit (tree search) |
+        | **Grounding** | Statistical text patterns | Causal dynamics |
+        
+        **Example:**
+        - **LLM**: "If I push a ball off a table..." → generates plausible *text*
+        - **World Model**: state(ball on table) + action(push) → predicts actual *trajectory*
+        
+        Language models learn *what sounds right*. World models learn *what actually happens*.
+        """)
+    
+    with gr.Accordion("🔬 Why does this matter for AI Safety?", open=False):
+        gr.Markdown("""
+        World models are crucial for AI safety research because:
+        
+        - **Predictability**: Agents that plan can be analyzed - we can inspect what futures they're considering
+        - **Corrigibility**: Planning agents can incorporate "avoid irreversible actions" into their search
+        - **Interpretability**: The model's predictions can be examined for accuracy and bias
+        - **Scalable Oversight**: Humans can audit the agent's "reasoning" by inspecting simulated futures
+        
+        Understanding how AI systems model the world helps us build systems we can trust and verify.
+        
+        ---
+        *Created by [Anthony Maio](https://huggingface.co/anthonym21) as an educational resource*
+        """)
+    
     # Connect buttons
     up_btn.click(lambda: do_action("up"), outputs=[grid_display, stats_display, message_display])
     down_btn.click(lambda: do_action("down"), outputs=[grid_display, stats_display, message_display])