Upload folder using huggingface_hub

chats/init context 10 AM.txt

@@ -1,685 +1,685 @@
-╭─── Claude Code v2.1.112 ─────────────────────────────────────────────────────╮
-│                                                    │ Tips for getting        │
-│                    Welcome back!                   │ started                 │
-│                                                    │ Run /init to create a … │
-│                       ▐▛███▜▌                      │ ─────────────────────── │
-│                      ▝▜█████▛▘                     │ Recent activity         │
-│                        ▘▘ ▝▝                       │ 2d ago  /init           │
-│                                                    │ /resume for more        │
-│           Sonnet 4.6 · API Usage Billing           │                         │
-│ ~\Projects\Hackathon\Meta PyTorch OpenEnv Hackath… │                         │
-╰──────────────────────────────────────────────────────────────────────────────╯
-
-  ↑ Install the PyCharm plugin from the JetBrains Marketplace: https://docs.claude.com/s/claude-code-jetbrains
-  ⎿  SessionStart:startup hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-❯ Read full codebase and make sure you don't make any changes to the code just  
-  read the code base.                                                           
-  ⎿  UserPromptSubmit hook error                                   ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1:
-
-● I'll explore the full codebase structure and read all the files.
-
-  Searched for 2 patterns (ctrl+o to expand)
-
-● Now I'll read all the source files. Let me start with the top-level files and
-  work through each directory.
-
-  Read 8 files (ctrl+o to expand)
-
-● Now reading all the Python source files.
-
-  Read 4 files (ctrl+o to expand)
-
-● Now reading the environment package files.
-
-  Read 5 files (ctrl+o to expand)
-
-● Now reading the scenarios, server, and test files.
-
-  Read 6 files (ctrl+o to expand)
-
-● Now reading the test files, scripts, and frontend assets.
-
-  Read 5 files (ctrl+o to expand)
-
-● Now reading the frontend static assets.
-
-  Read 5 files (ctrl+o to expand)
-
-● Now reading the CSS files and remaining static assets.
-
-  Read 5 files (ctrl+o to expand)
-
-● Finally, reading the egg-info metadata files.
-
-  Read 5 files (ctrl+o to expand)
-
-● I've read the entire codebase. Here's a complete summary:
-
-  ---
-  MaterialForge - Full Codebase Summary
-
-  Project: A reinforcement learning environment for autonomous crystalline
-  structure discovery, built for the Meta PyTorch OpenEnv Hackathon. Agents
-  place atoms (Metal/A, Conductor/B, Ceramic/C, Polymer/P) on an 8x8 lattice to
-  match target physical properties (Hardness, Conductivity, Thermal Resistance,
-  Elasticity).
-
-  Architecture
-
-  MaterialForge/
-  ├── environment/          # Core physics engine
-  │   ├── config.py         # Constants: grid size (8), atom types, costs,
-  difficulty presets
-  │   ├── lattice.py        # Lattice class: 8x8 grid with place/replace/remove
-  operations
-  │   ├── physics.py        # Property estimation, phase classification,
-  stability/order scoring
-  │   └── rubrics.py        # HeuristicRewardRubric: 50% property match, 25%
-  stability, 15% order, 10% phase
-  ├── server/               # FastAPI server (OpenEnv-compliant)
-  │   ├── app.py            # Entry point: creates FastAPI app, mounts static
-  files, /playground route
-  │   ├── material_forge_env_environment.py  # MaterialForgeEnvironment:
-  reset/step logic, done detection
-  │   └── static/           # "Discovery Lab" dashboard (HTML/CSS/JS with
-  Chart.js)
-  ├── scenarios/
-  │   └── scenarios.py      # 5 named scenarios (diamond-like, conductor,
-  heat-shield, flexible-polymer, balanced-alloy)
-  ├── models.py             # Pydantic models: MaterialForgeAction,
-  MaterialForgeObservation
-  ├── client.py             # WebSocket client (EnvClient subclass)
-  ├── inference.py          # Agent: heuristic lookahead + LLM
-  (OpenAI-compatible) for action selection
-  ├── tests/
-  │   └── benchmark.py      # Analytical benchmark suite (100-trial evaluations,
-   Ollama support)
-  ├── openenv.yaml          # 3 tasks: basic-synthesis, diamond-like,
-  superconductor-analogue
-  ├── Dockerfile            # Multi-stage Docker build for HF Spaces deployment
-  └── pyproject.toml        # Python 3.10+, depends on openenv-core, openai,
-  gradio
-
-  Key Components
-
-  - Physics engine (physics.py): Calculates properties from atom fractions +
-  neighbor bonding + BFS percolation (for conductivity). Classifies phase via
-  2x2 pattern repetition analysis. Computes structural stability (coordination
-  energy + mirror symmetry) and lattice order (Bragg-like + quadrant
-  homogeneity).
-  - Inference agent (inference.py): Generates candidate actions via heuristic
-  scoring (clustering, center placement), simulates each locally via the physics
-   engine, then asks an LLM to pick the best from top-8 candidates. Falls back
-  to pure heuristic if LLM fails.
-  - Frontend dashboard: Dark-themed professional UI with real-time lattice
-  canvas, property alignment bars, reward convergence chart, KPIs, and analytics
-   tab. WebSocket-connected to the server.
-  - 3 difficulty levels: Easy (tolerance 20, budget 120), Medium (tolerance 10,
-  budget 80), Hard (tolerance 5, budget 60).
-
-  No code changes were made.
-
-● Ran 1 stop hook (ctrl+o to expand)
-  ⎿  Stop hook error: Failed with non-blocking status code:
-  /usr/bin/bash: line 1: C:Program: command not found
-
-✻ Worked for 7m 27s
-
-❯ We have already qualified Round 1 and are now in the Grand Finale. Our        
-  environment is largely complete and functional. The remaining work likely     
-  involves training, final packaging, and required submissions.                 
-                                                                                
-  Your task is to determine EXACTLY what actions are required next.             
-                                                                                
-  You must carefully analyze the following resources:                           
-                                                                                
-  VIDEO                                                                         
-  https://youtu.be/Jew4lhAiqnw?si=BzcMT5UkCfjtwvp-                              
-                                                                                
-  DOCUMENTS                                                                     
-  https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlEu8M9Z 
-  GU/edit?tab=t.0                                                               
-  https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGcxbC6m 
-  Fs/preview?pru=AAABnebPvLA*TA1iaP5LykX7FUN4W27LdQ&tab=t.0#heading=h.b0qzylowp 
-  ixh                                                                           
-  https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11ERvgg 
-  PM/edit?tab=t.0                                                               
-  https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb0wDj2 
-  o4/edit?tab=t.0#heading=h.iksn3i4lpwnt                                        
-  https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGcxbC6m 
-  Fs/edit?usp=drivesdk                                                          
-                                                                                
-  Instructions:                                                                 
-                                                                                
-  1. Read every document completely.                                            
-  2. Visit every link referenced inside those documents.                        
-  3. Extract all instructions, requirements, deadlines, and submission steps    
-  related to the FINAL ROUND.                                                   
-  4. Watch the entire video and extract any guidance or rules mentioned by the  
-  organizers.                                                                   
-  5. Identify any required artifacts or deliverables (examples: environment     
-  repo, training results, evaluation scripts, demos, presentations, etc.).      
-                                                                                
-  After reviewing everything, produce the following structured output:          
-                                                                                
-  SECTION 1 — Final Round Requirements                                          
-  List all explicit requirements given by the organizers.                       
-                                                                                
-  SECTION 2 — Required Deliverables                                             
-  List everything the team must submit or demonstrate.                          
-                                                                                
-  SECTION 3 — Training Requirements                                             
-  Identify:                                                                     
-  - whether RL training is expected                                             
-  - recommended algorithms or tooling                                           
-  - expected evaluation metrics or benchmarks                                   
-                                                                                
-  SECTION 4 — Submission Process                                                
-  Explain:                                                                      
-  - where submissions happen                                                    
-  - what files or links are needed                                              
-  - any format requirements.                                                    
-                                                                                
-  SECTION 5 — Judging Criteria                                                  
-  Extract how projects will be evaluated.                                       
-                                                                                
-  SECTION 6 — Immediate Action Plan                                             
-  Create a prioritized checklist of the next steps our team should take to      
-  maximize our chances of winning.                                              
-                                                                                
-  Focus on actionable insights only. Do not summarize the documents broadly.    
-  Extract only the information relevant to what we must do next. Dont make any  
-  code changes yet                                                              
-  ⎿  UserPromptSubmit hook error                                   ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1:
-
-● I'll fetch all the documents and video to extract the finale requirements.
-  Starting with all resources in parallel.
-
-● Fetch(https://youtu.be/Jew4lhAiqnw?si=BzcMT5UkCfjtwvp-)
-  ⎿  Error: Request failed with status code 303
-
-● Fetch(https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlE
-       u8M9ZGU/edit?tab=t.0)
-  ⎿  Received 574.3KB (200 OK)
-
-● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
-       xbC6mFs/preview?pru=AAABnebPvLA*TA1iaP5LykX7FUN4W27LdQ&tab=t.0#heading=h.
-       b0qzylowpixh)
-  ⎿  Received 282KB (200 OK)
-
-● Fetch(https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11
-       ERvggPM/edit?tab=t.0)
-  ⎿  Received 198.3KB (200 OK)
-
-● Fetch(https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb
-       0wDj2o4/edit?tab=t.0#heading=h.iksn3i4lpwnt)  ⎿  Received 519.4KB (200 OK)
-
-● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
-       xbC6mFs/edit?usp=drivesdk)
-  ⎿  Received 358.2KB (200 OK)
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-● The Google Docs are behind authentication and the YouTube video redirected.
-  Let me try alternative approaches to access these resources.
-
-● Fetch(https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlE
-       u8M9ZGU/export?format=txt)
-  ⎿  Received 820 bytes (307 Temporary Redirect)
-
-● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
-       xbC6mFs/export?format=txt)
-  ⎿  Received 820 bytes (307 Temporary Redirect)
-
-● Fetch(https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11
-       ERvggPM/export?format=txt)
-  ⎿  Received 820 bytes (307 Temporary Redirect)
-
-● Fetch(https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb
-       0wDj2o4/export?format=txt)
-  ⎿  Received 820 bytes (307 Temporary Redirect)
-
-● Fetch(https://www.youtube.com/watch?v=Jew4lhAiqnw)
-  ⎿  Received 1.1MB (200 OK)
-
-● Following the redirects to get the actual document content.
-
-● Fetch(https://doc-08-8o-docstext.googleusercontent.com/export/e6hpso97lrhpva19
-       l3uocm525o/mesgsm24tv6fcjskeodkl8v654/1777090195000/110811587972687932588
-       /*/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlEu8M9ZGU?format=txt)
-  ⎿  Error: Request failed with status code 400
-
-● Fetch(https://doc-14-bs-docstext.googleusercontent.com/export/e6hpso97lrhpva19
-       l3uocm525o/udraaum8cqmoegp5o6gm5f37rs/1777090195000/107073175634533162618
-       /*/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGcxbC6mFs?format=txt)
-  ⎿  Error: Request failed with status code 400
-
-● Fetch(https://doc-0k-40-docstext.googleusercontent.com/export/e6hpso97lrhpva19
-       l3uocm525o/igqp5g3rqgl01k2sbhph5jasd8/1777090195000/104102298771320400411
-       /*/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11ERvggPM?format=txt)
-  ⎿  Error: Request failed with status code 400
-
-● Fetch(https://doc-0k-bs-docstext.googleusercontent.com/export/e6hpso97lrhpva19
-       l3uocm525o/50ntp2srn6mv6e0t942chcuo4s/1777090200000/107073175634533162618
-       /*/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb0wDj2o4?format=txt)
-  ⎿  Error: Request failed with status code 400
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-● The redirects failed due to the wildcard in the URL. Let me try the
-  published/public versions and also search for the hackathon details online.
-
-● Fetch(https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlE
-       u8M9ZGU/pub)
-  ⎿  Error: Request failed with status code 401
-
-● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
-       xbC6mFs/pub)
-  ⎿  Error: Request failed with status code 401
-
-● Fetch(https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11
-       ERvggPM/pub)
-  ⎿  Error: Request failed with status code 401
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-● Fetch(https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb
-       0wDj2o4/pub)
-  ⎿  Error: Request failed with status code 401
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-● Web Search("Meta PyTorch OpenEnv Hackathon 2026 Scaler finale judging criteria
-             submission requirements")
-  ⎿  Did 10 searches in 8s
-  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-● Fetch(https://dev.to/alara_iris_5d51447594b8a3/meta-pytorch-openenv-hackathon-
-       x-sst-10nd)
-  ⎿  Received 72.1KB (200 OK)
-
-● Fetch(https://www.scaler.com/school-of-technology/meta-pytorch-hackathon)
-  ⎿  Received 8.5KB (200 OK)
-
-● Fetch(https://kalinga.ai/openenv-ai-hackathon-india-2026/)
-  ⎿  Received 193KB (200 OK)
-
-● Web Search("OpenEnv hackathon round 2 finale requirements training GRPO
-            submission themes judging 2026")
-  ⎿  Did 10 searches in 7s
-
-● Fetch(https://github.com/meta-pytorch/OpenEnv)
-  ⎿  Received 434.6KB (200 OK)
-
-● Web Search(""openenv" hackathon themes "world modeling" "long horizon" judging
-             criteria scoring environment 2026")
-  ⎿  Did 8 searches in 10s
-
-● Web Search("OpenEnv hackathon "round 2" OR "finale" training pipeline GRPO
-            unsloth TRL submission requirements environment")
-  ⎿  Did 6 searches in 11s
-
-● Fetch(https://pytorch.org/event/openenv-ai-hackathon/)
-  ⎿  Received 160.2KB (200 OK)
-
-● Fetch(https://huggingface.co/docs/trl/openenv)
-  ⎿  Received 27.2KB (200 OK)
-
-● Fetch(https://deepwiki.com/meta-pytorch/OpenEnv/12-examples-and-tutorials)
-  ⎿  Interrupted · What should Claude do instead?
-
-❯ If you are unable to read via url I'll provide you the raw text               
-  ⎿  UserPromptSubmit hook error                                       ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
-
-● Yes, please paste the raw text from the documents. I've gathered significant
-  context from public sources (Scaler, PyTorch.org, TRL docs, GitHub), but the
-  Google Docs are behind authentication and contain the critical finale-specific
-   details I need — especially the Themes & Judging Criteria document and the
-  Participant Help Guide.
-
-  Please paste the content from these documents:
-
-  1. Hackathon FAQs (1ysO...)
-  2. Themes & Judging Criteria (1Odz...)
-  3. OpenEnv Hackathon Resources (1m4E...)
-  4. Participant Help Guide (1lGO...)
-
-  Also, if you can share the key points from the YouTube video (Mega Lecture
-  91), that would fill the remaining gaps. I'll compile the full action plan
-  once I have the raw text.
-
-● Ran 1 stop hook (ctrl+o to expand)
-  ⎿  Stop hook error: Failed with non-blocking status code:
-  /usr/bin/bash: line 1: C:Program: command not found
-
-❯ I have added everything in @docs\                                             
-  ⎿  Listed directory docs\
-  ⎿  UserPromptSubmit hook error                                       ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
-
-● The documents are in the repo. Let me read all of them now.
-
-  Read 6 files (ctrl+o to expand)
-  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
-  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
-  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
-  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
-
-● The Themes & Judging Criteria doc and the FAQ are large. Let me read them in
-  parts.
-
-  Searched for 1 pattern, read 2 files, listed 1 directory (ctrl+o to expand)  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:Bash hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-  ⎿  PostToolUseFailure:Bash hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-● Now reading the FAQ and the transcript (key sections).
-
-  Searched for 1 pattern, read 1 file (ctrl+o to expand)
-  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
-
-● Now let me read the key portions of the transcript for hackathon-specific
-  guidance.
-
-  Searched for 1 pattern (ctrl+o to expand)
-
-● I now have all the critical information. Let me also verify the Themes
-  document is fully captured since I read it via bash.
-
-  Searched for 1 pattern, read 1 file (ctrl+o to expand)
-
-● I now have all the information needed. Here is the complete analysis:
-
-  ---
-  SECTION 1 — Final Round Requirements
-
-  Explicit organizer requirements:
-
-  1. Use OpenEnv (latest release) — build on top of the framework, don't
-  reinvent the wheel
-  2. A working training script using Unsloth or HF TRL, ideally as a Colab 
-  notebook so judges can re-run it
-  3. Evidence of actual training — at minimum, loss and reward plots from a real
-   run
-  4. A short writeup: mini-blog on HuggingFace OR < 2 minute video on YouTube,
-  OR a short slide deck
-  5. Environment hosted on Hugging Face Spaces — discoverable and runnable;
-  judges will pull the environment from the submitted URL
-  6. A README that motivates the problem, explains the environment, and shows
-  results — with links to HF Space, blog/video, and all supplementary materials
-  7. One submission per team — pick your best idea
-  8. No large video files in the HF Hub submission — use URL references instead
-  9. No changes/commits after the submission deadline — judges evaluate the
-  snapshot at deadline
-  10. It is NOT mandatory to keep the same problem statement as Round 1 — only
-  keep it if it aligns with the finale themes
-  11. Compute credits for HuggingFace will be given onsite on April 25-26 for
-  post-training
-
-  ---
-  SECTION 2 — Required Deliverables
-
-  #: 1
-  Deliverable: OpenEnv-compliant environment
-  Format: HF Space (live URL)
-  Status for MaterialForge: Partially done — deployed at
-    ArshPathan/material_forge_env but needs verification with latest OpenEnv
-  ────────────────────────────────────────
-  #: 2
-  Deliverable: Training script
-  Format: Colab notebook (Unsloth or TRL + GRPO)
-  Status for MaterialForge: NOT DONE — critical gap
-  ────────────────────────────────────────
-  #: 3
-  Deliverable: Training evidence
-  Format: Reward curves, loss plots, before/after comparison (.png/.jpg in repo
-  +
-    README)
-  Status for MaterialForge: NOT DONE — critical gap
-  ────────────────────────────────────────
-  #: 4
-  Deliverable: Mini-blog OR <2min video OR slide deck
-  Format: HF blog post / YouTube / Slides
-  Status for MaterialForge: NOT DONE
-  ────────────────────────────────────────
-  #: 5
-  Deliverable: README with results
-  Format: Markdown in repo root
-  Status for MaterialForge: Exists but needs major update with training results,
-
-    plots, and links
-  ────────────────────────────────────────
-  #: 6
-  Deliverable: openenv.yaml manifest
-  Format: YAML in repo root
-  Status for MaterialForge: Done
-  ────────────────────────────────────────
-  #: 7
-  Deliverable: Valid client/server separation
-  Format: Code architecture
-  Status for MaterialForge: Done
-  ────────────────────────────────────────
-  #: 8
-  Deliverable: 3-minute live pitch
-  Format: In-person presentation
-  Status for MaterialForge: NOT PREPARED
-
-  ---
-  SECTION 3 — Training Requirements
-
-  Is RL training expected?
-
-  YES — absolutely mandatory. The documents are explicit:
-
-  ▎ "The bar isn't 'training script exists.' The bar is 'training script runs 
-  ▎ against the environment, the agent learns, and you can show it.'"
-
-  Training evidence is worth 20% of the score directly, and the training
-  pipeline is another 10%.
-
-  Recommended stack:
-
-  - TRL (GRPOTrainer) — the primary training framework
-  - Unsloth — for memory efficiency (QLoRA/LoRA), especially on limited hardware
-  - GRPO algorithm — preferred over PPO for LLM RL; more memory-efficient, no
-  value model needed
-  - OpenEnv environment_factory pattern — wraps your environment as tool methods
-   the trainer discovers automatically
-
-  How training should work (from TRL docs):
-
-  1. Create an environment class with reset() and tool methods (e.g.,
-  place_atom(), remove_atom())
-  2. The GRPOTrainer discovers tool methods, generates completions, parses tool
-  calls, executes them, feeds results back
-  3. Reward comes from the environment after each episode
-  4. Model: Start with a capable instruct model (e.g., Qwen/Qwen3-0.6B or
-  Qwen/Qwen2.5-3B)
-
-  Expected evaluation metrics:
-
-  - Reward curves over training steps (must be plotted, labeled, saved as
-  images)
-  - Before/after behavior comparison (untrained baseline vs. trained agent)
-  - Loss plots
-  - Quantitative comparison: random agent vs. heuristic baseline vs. trained LLM
-   agent
-
-  Key warnings from organizers:
-
-  - Don't naively upcast 4-bit to 16-bit when saving LoRA — use proper merge
-  path
-  - Start with easy tasks, use curriculum learning
-  - Binary rewards (1.0 success / 0.0 fail) often work better than shaped
-  rewards for GRPO
-  - max_completion_length must be increased for long episodes (default is too
-  small)
-  - Test reward function before training — if a capable model can't beat random,
-   reward is broken
-
-  ---
-  SECTION 4 — Submission Process
-
-  Where submissions happen:
-
-  - HuggingFace Spaces — the environment must be live and accessible via URL
-  - Judges pull the environment from the submitted URL to evaluate it
-  - All supplementary materials (blog, video, slides, plots) must be linked from
-   the README
-
-  What files/links are needed:
-
-  1. HF Space URL of the running environment
-  2. Colab notebook with training script (linked from README)
-  3. Plot images (.png/.jpg) committed to the repo and embedded in README
-  4. Blog/video/slides link in README
-  5. openenv.yaml manifest in repo root
-  6. README.md with: problem motivation, how env works, results, and all links
-
-  Format requirements:
-
-  - Plots: labeled axes, .png or .jpg, embedded in README with one-line captions
-  - Video: < 2 minutes
-  - No large binary files in HF Hub
-  - README readable in 3-5 minutes
-
-  ---
-  SECTION 5 — Judging Criteria
-
-  Criterion: Environment Innovation
-  Weight: 40%
-  What judges evaluate: Is it novel, creative, genuinely challenging? Does it
-    teach an LLM something it can't currently do well? Underexplored domain?
-    Fresh angle? (Judges have seen too many chess/snake/grid-world clones)
-  ────────────────────────────────────────
-  Criterion: Storytelling & Presentation
-  Weight: 30%
-  What judges evaluate: Clear explanation of problem, environment, and what the
-    agent learned. Engaging demo, easy to follow for non-technical audience.
-    README answers: Problem → Environment → Results → Why it matters
-  ────────────────────────────────────────
-  Criterion: Showing Improvement in Rewards
-  Weight: 20%
-  What judges evaluate: Observable training progress: reward curves,
-  before/after
-    behavior, baseline comparison. Proves the agent actually learned something
-  ────────────────────────────────────────
-  Criterion: Reward & Training Pipeline
-  Weight: 10%
-  What judges evaluate: Coherent reward logic, pipeline produces meaningful
-    improvement in agent behavior
-
-  What makes a submission stand out:
-
-  - Ambitious, original problem — "Does this environment exist to teach an LLM
-  something it currently can't do well?"
-  - Rich reward signal — not just 0/1, uses OpenEnv's Rubric system, hard to
-  game
-  - Real end-to-end training — not just "script exists" but "script runs, agent
-  learns, you can show it"
-  - Readable plots on same axes for comparison
-  - A story, not an API doc — problem gap → environment → results → why it
-  matters
-
-  MaterialForge's theme alignment:
-
-  MaterialForge fits Theme #3.1 (World Modeling — Professional Tasks) and Theme 
-  #2 (Long-Horizon Planning). The scientific workflow (inverse molecular design)
-   is a genuinely novel domain for LLM RL training — this is a strong
-  differentiator.
-
-  ---
-  SECTION 6 — Immediate Action Plan
-
-  CRITICAL PATH (must-do, in priority order):
-
-  - 1. Build the TRL/GRPO Training Script (HIGHEST PRIORITY)
-    - Create a Colab notebook that uses GRPOTrainer with environment_factory
-    - Wrap MaterialForge as a TRL-compatible environment class with tool methods
-   (place_atom, remove_atom, replace_atom) instead of a generic step()
-    - Use Unsloth for memory efficiency
-    - Start model: Qwen/Qwen3-0.6B or Qwen/Qwen2.5-3B
-    - Connect to the MaterialForge HF Space via WebSocket
-    - This unlocks 30% of the score (20% improvement evidence + 10% pipeline)
-  - 2. Run Training & Capture Evidence
-    - Execute training with at least enough steps for curves to be meaningful
-    - Save reward curves and loss plots as .png with labeled axes
-    - Record baseline comparison: random agent score vs. heuristic vs. trained
-  LLM
-    - Use the onsite HF compute credits (April 25-26) for larger runs
-    - Before/after qualitative examples (grid screenshots showing trained agent
-  building better crystals)
-  - 3. Update README with Results
-    - Restructure to tell the story: Problem → Environment → Results → Why It
-  Matters
-    - Embed training plots with one-line captions
-    - Add links to: HF Space, Colab training notebook, video/blog
-    - Show baseline vs. trained comparison table
-    - Make it readable in 3-5 minutes
-  - 4. Create Mini-Blog or <2min Video
-    - Option A: HuggingFace blog post explaining the environment and training
-  results
-    - Option B: YouTube video (<2 min) showing the Discovery Lab dashboard +
-  training progress
-    - Option C: Short slide deck for the pitch
-    - Link it from README
-  - 5. Verify HF Space is Live & Current
-    - Ensure ArshPathan/material_forge_env is running and accessible
-    - Test that judges can hit /reset and /step endpoints
-    - Verify openenv.yaml is valid with latest OpenEnv release
-  - 6. Prepare 3-Minute Pitch
-    - Structure: Problem (inverse molecular design) → Environment (what agent
-  sees/does) → Demo (show dashboard) → Training Results (reward curves,
-  before/after) → Why it matters (scientific discovery)
-    - Practice delivering in 3 minutes
-    - Emphasize the novelty: scientific professional workflow, not another game
-  clone
-
-  NICE-TO-HAVE (if time permits):
-
-  - Add curriculum learning (easy → medium → hard scenarios during training)
-  - Multiple reward function ablations to show robustness
-  - Anti-reward-hacking analysis
-  - WandB integration for richer training logs
-
-  Time allocation suggestion (for the 48-hour finale):
-
-  ┌───────────┬───────┬─────────────────────────────────────────────────────┐
-  │   Phase   │ Hours │                      Activity                       │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours 0-4 │ 4h    │ Build TRL training script + environment wrapper in  │
-  │           │       │ Colab                                               │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours 4-8 │ 4h    │ Debug training loop, get first reward curves        │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours     │ 8h    │ Run training with onsite compute credits, iterate   │
-  │ 8-16      │       │ on rewards                                          │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours     │ 8h    │ Capture all evidence, generate plots, run baseline  │
-  │ 16-24     │       │ comparisons                                         │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours     │ 8h    │ Update README, write blog/create video, clean up    │
-  │ 24-32     │       │ repo                                                │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours     │ 8h    │ Polish demo, prepare pitch, verify HF Space         │
-  │ 32-40     │       │                                                     │
-  ├───────────┼───────┼─────────────────────────────────────────────────────┤
-  │ Hours     │ 8h    │ Final submission prep, rehearse pitch, buffer for   │
-  │ 40-48     │       │ issues                                              │
-  └───────────┴───────┴─────────────────────────────────────────────────────┘
-
-● User answered Claude's questions:
-  ⎿  · [Bonsai] What went well or could be improved on this task? Consider 
-     correctness, performance, and code quality. → Awsome performance
-
-● User answered Claude's questions:
-  ⎿  · [Bonsai] How did the model do on this task? → Great
-
-● Thank you for the feedback!
-
-● Ran 1 stop hook (ctrl+o to expand)
-  ⎿  Stop hook error: Failed with non-blocking status code:
-  /usr/bin/bash: line 1: C:Program: command not found
-
+╭─── Claude Code v2.1.112 ─────────────────────────────────────────────────────╮
+│                                                    │ Tips for getting        │
+│                    Welcome back!                   │ started                 │
+│                                                    │ Run /init to create a … │
+│                       ▐▛███▜▌                      │ ─────────────────────── │
+│                      ▝▜█████▛▘                     │ Recent activity         │
+│                        ▘▘ ▝▝                       │ 2d ago  /init           │
+│                                                    │ /resume for more        │
+│           Sonnet 4.6 · API Usage Billing           │                         │
+│ ~\Projects\Hackathon\Meta PyTorch OpenEnv Hackath… │                         │
+╰──────────────────────────────────────────────────────────────────────────────╯
+
+  ↑ Install the PyCharm plugin from the JetBrains Marketplace: https://docs.claude.com/s/claude-code-jetbrains
+  ⎿  SessionStart:startup hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+❯ Read full codebase and make sure you don't make any changes to the code just  
+  read the code base.                                                           
+  ⎿  UserPromptSubmit hook error                                   ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1:
+
+● I'll explore the full codebase structure and read all the files.
+
+  Searched for 2 patterns (ctrl+o to expand)
+
+● Now I'll read all the source files. Let me start with the top-level files and
+  work through each directory.
+
+  Read 8 files (ctrl+o to expand)
+
+● Now reading all the Python source files.
+
+  Read 4 files (ctrl+o to expand)
+
+● Now reading the environment package files.
+
+  Read 5 files (ctrl+o to expand)
+
+● Now reading the scenarios, server, and test files.
+
+  Read 6 files (ctrl+o to expand)
+
+● Now reading the test files, scripts, and frontend assets.
+
+  Read 5 files (ctrl+o to expand)
+
+● Now reading the frontend static assets.
+
+  Read 5 files (ctrl+o to expand)
+
+● Now reading the CSS files and remaining static assets.
+
+  Read 5 files (ctrl+o to expand)
+
+● Finally, reading the egg-info metadata files.
+
+  Read 5 files (ctrl+o to expand)
+
+● I've read the entire codebase. Here's a complete summary:
+
+  ---
+  MaterialForge - Full Codebase Summary
+
+  Project: A reinforcement learning environment for autonomous crystalline
+  structure discovery, built for the Meta PyTorch OpenEnv Hackathon. Agents
+  place atoms (Metal/A, Conductor/B, Ceramic/C, Polymer/P) on an 8x8 lattice to
+  match target physical properties (Hardness, Conductivity, Thermal Resistance,
+  Elasticity).
+
+  Architecture
+
+  MaterialForge/
+  ├── environment/          # Core physics engine
+  │   ├── config.py         # Constants: grid size (8), atom types, costs,
+  difficulty presets
+  │   ├── lattice.py        # Lattice class: 8x8 grid with place/replace/remove
+  operations
+  │   ├── physics.py        # Property estimation, phase classification,
+  stability/order scoring
+  ��   └── rubrics.py        # HeuristicRewardRubric: 50% property match, 25%
+  stability, 15% order, 10% phase
+  ├── server/               # FastAPI server (OpenEnv-compliant)
+  │   ├── app.py            # Entry point: creates FastAPI app, mounts static
+  files, /playground route
+  │   ├── material_forge_env_environment.py  # MaterialForgeEnvironment:
+  reset/step logic, done detection
+  │   └── static/           # "Discovery Lab" dashboard (HTML/CSS/JS with
+  Chart.js)
+  ├── scenarios/
+  │   └── scenarios.py      # 5 named scenarios (diamond-like, conductor,
+  heat-shield, flexible-polymer, balanced-alloy)
+  ├── models.py             # Pydantic models: MaterialForgeAction,
+  MaterialForgeObservation
+  ├── client.py             # WebSocket client (EnvClient subclass)
+  ├── inference.py          # Agent: heuristic lookahead + LLM
+  (OpenAI-compatible) for action selection
+  ├── tests/
+  │   └── benchmark.py      # Analytical benchmark suite (100-trial evaluations,
+   Ollama support)
+  ├── openenv.yaml          # 3 tasks: basic-synthesis, diamond-like,
+  superconductor-analogue
+  ├── Dockerfile            # Multi-stage Docker build for HF Spaces deployment
+  └── pyproject.toml        # Python 3.10+, depends on openenv-core, openai,
+  gradio
+
+  Key Components
+
+  - Physics engine (physics.py): Calculates properties from atom fractions +
+  neighbor bonding + BFS percolation (for conductivity). Classifies phase via
+  2x2 pattern repetition analysis. Computes structural stability (coordination
+  energy + mirror symmetry) and lattice order (Bragg-like + quadrant
+  homogeneity).
+  - Inference agent (inference.py): Generates candidate actions via heuristic
+  scoring (clustering, center placement), simulates each locally via the physics
+   engine, then asks an LLM to pick the best from top-8 candidates. Falls back
+  to pure heuristic if LLM fails.
+  - Frontend dashboard: Dark-themed professional UI with real-time lattice
+  canvas, property alignment bars, reward convergence chart, KPIs, and analytics
+   tab. WebSocket-connected to the server.
+  - 3 difficulty levels: Easy (tolerance 20, budget 120), Medium (tolerance 10,
+  budget 80), Hard (tolerance 5, budget 60).
+
+  No code changes were made.
+
+● Ran 1 stop hook (ctrl+o to expand)
+  ⎿  Stop hook error: Failed with non-blocking status code:
+  /usr/bin/bash: line 1: C:Program: command not found
+
+✻ Worked for 7m 27s
+
+❯ We have already qualified Round 1 and are now in the Grand Finale. Our        
+  environment is largely complete and functional. The remaining work likely     
+  involves training, final packaging, and required submissions.                 
+                                                                                
+  Your task is to determine EXACTLY what actions are required next.             
+                                                                                
+  You must carefully analyze the following resources:                           
+                                                                                
+  VIDEO                                                                         
+  https://youtu.be/Jew4lhAiqnw?si=BzcMT5UkCfjtwvp-                              
+                                                                                
+  DOCUMENTS                                                                     
+  https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlEu8M9Z 
+  GU/edit?tab=t.0                                                               
+  https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGcxbC6m 
+  Fs/preview?pru=AAABnebPvLA*TA1iaP5LykX7FUN4W27LdQ&tab=t.0#heading=h.b0qzylowp 
+  ixh                                                                           
+  https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11ERvgg 
+  PM/edit?tab=t.0                                                               
+  https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb0wDj2 
+  o4/edit?tab=t.0#heading=h.iksn3i4lpwnt                                        
+  https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGcxbC6m 
+  Fs/edit?usp=drivesdk                                                          
+                                                                                
+  Instructions:                                                                 
+                                                                                
+  1. Read every document completely.                                            
+  2. Visit every link referenced inside those documents.                        
+  3. Extract all instructions, requirements, deadlines, and submission steps    
+  related to the FINAL ROUND.                                                   
+  4. Watch the entire video and extract any guidance or rules mentioned by the  
+  organizers.                                                                   
+  5. Identify any required artifacts or deliverables (examples: environment     
+  repo, training results, evaluation scripts, demos, presentations, etc.).      
+                                                                                
+  After reviewing everything, produce the following structured output:          
+                                                                                
+  SECTION 1 — Final Round Requirements                                          
+  List all explicit requirements given by the organizers.                       
+                                                                                
+  SECTION 2 — Required Deliverables                                             
+  List everything the team must submit or demonstrate.                          
+                                                                                
+  SECTION 3 — Training Requirements                                             
+  Identify:                                                                     
+  - whether RL training is expected                                             
+  - recommended algorithms or tooling                                           
+  - expected evaluation metrics or benchmarks                                   
+                                                                                
+  SECTION 4 — Submission Process                                                
+  Explain:                                                                      
+  - where submissions happen                                                    
+  - what files or links are needed                                              
+  - any format requirements.                                                    
+                                                                                
+  SECTION 5 — Judging Criteria                                                  
+  Extract how projects will be evaluated.                                       
+                                                                                
+  SECTION 6 — Immediate Action Plan                                             
+  Create a prioritized checklist of the next steps our team should take to      
+  maximize our chances of winning.                                              
+                                                                                
+  Focus on actionable insights only. Do not summarize the documents broadly.    
+  Extract only the information relevant to what we must do next. Dont make any  
+  code changes yet                                                              
+  ⎿  UserPromptSubmit hook error                                   ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1:
+
+● I'll fetch all the documents and video to extract the finale requirements.
+  Starting with all resources in parallel.
+
+● Fetch(https://youtu.be/Jew4lhAiqnw?si=BzcMT5UkCfjtwvp-)
+  ⎿  Error: Request failed with status code 303
+
+● Fetch(https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlE
+       u8M9ZGU/edit?tab=t.0)
+  ⎿  Received 574.3KB (200 OK)
+
+● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
+       xbC6mFs/preview?pru=AAABnebPvLA*TA1iaP5LykX7FUN4W27LdQ&tab=t.0#heading=h.
+       b0qzylowpixh)
+  ⎿  Received 282KB (200 OK)
+
+● Fetch(https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11
+       ERvggPM/edit?tab=t.0)
+  ⎿  Received 198.3KB (200 OK)
+
+● Fetch(https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb
+       0wDj2o4/edit?tab=t.0#heading=h.iksn3i4lpwnt)  ⎿  Received 519.4KB (200 OK)
+
+● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
+       xbC6mFs/edit?usp=drivesdk)
+  ⎿  Received 358.2KB (200 OK)
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+● The Google Docs are behind authentication and the YouTube video redirected.
+  Let me try alternative approaches to access these resources.
+
+● Fetch(https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlE
+       u8M9ZGU/export?format=txt)
+  ⎿  Received 820 bytes (307 Temporary Redirect)
+
+● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
+       xbC6mFs/export?format=txt)
+  ⎿  Received 820 bytes (307 Temporary Redirect)
+
+● Fetch(https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11
+       ERvggPM/export?format=txt)
+  ⎿  Received 820 bytes (307 Temporary Redirect)
+
+● Fetch(https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb
+       0wDj2o4/export?format=txt)
+  ⎿  Received 820 bytes (307 Temporary Redirect)
+
+● Fetch(https://www.youtube.com/watch?v=Jew4lhAiqnw)
+  ⎿  Received 1.1MB (200 OK)
+
+● Following the redirects to get the actual document content.
+
+● Fetch(https://doc-08-8o-docstext.googleusercontent.com/export/e6hpso97lrhpva19
+       l3uocm525o/mesgsm24tv6fcjskeodkl8v654/1777090195000/110811587972687932588
+       /*/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlEu8M9ZGU?format=txt)
+  ⎿  Error: Request failed with status code 400
+
+● Fetch(https://doc-14-bs-docstext.googleusercontent.com/export/e6hpso97lrhpva19
+       l3uocm525o/udraaum8cqmoegp5o6gm5f37rs/1777090195000/107073175634533162618
+       /*/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGcxbC6mFs?format=txt)
+  ⎿  Error: Request failed with status code 400
+
+● Fetch(https://doc-0k-40-docstext.googleusercontent.com/export/e6hpso97lrhpva19
+       l3uocm525o/igqp5g3rqgl01k2sbhph5jasd8/1777090195000/104102298771320400411
+       /*/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11ERvggPM?format=txt)
+  ⎿  Error: Request failed with status code 400
+
+● Fetch(https://doc-0k-bs-docstext.googleusercontent.com/export/e6hpso97lrhpva19
+       l3uocm525o/50ntp2srn6mv6e0t942chcuo4s/1777090200000/107073175634533162618
+       /*/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb0wDj2o4?format=txt)
+  ⎿  Error: Request failed with status code 400
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+● The redirects failed due to the wildcard in the URL. Let me try the
+  published/public versions and also search for the hackathon details online.
+
+● Fetch(https://docs.google.com/document/d/1ysOVuq6StZpJ7mcyNi4bpmxxB1AnAWE9YGlE
+       u8M9ZGU/pub)
+  ⎿  Error: Request failed with status code 401
+
+● Fetch(https://docs.google.com/document/d/1Odznuzwtb1ecDOm2t6ToZd4MuMXXfO6vWUGc
+       xbC6mFs/pub)
+  ⎿  Error: Request failed with status code 401
+
+● Fetch(https://docs.google.com/document/d/1m4EEBkomp3UJ4Wc-ydRjITDfdzLWPyFqNx11
+       ERvggPM/pub)
+  ⎿  Error: Request failed with status code 401
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+● Fetch(https://docs.google.com/document/d/1lGOb5w8LGSbYd7_L-L-zAD9moThrMQZpUyIb
+       0wDj2o4/pub)
+  ⎿  Error: Request failed with status code 401
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+● Web Search("Meta PyTorch OpenEnv Hackathon 2026 Scaler finale judging criteria
+             submission requirements")
+  ⎿  Did 10 searches in 8s
+  ⎿  PostToolUseFailure:WebFetch hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+● Fetch(https://dev.to/alara_iris_5d51447594b8a3/meta-pytorch-openenv-hackathon-
+       x-sst-10nd)
+  ⎿  Received 72.1KB (200 OK)
+
+● Fetch(https://www.scaler.com/school-of-technology/meta-pytorch-hackathon)
+  ⎿  Received 8.5KB (200 OK)
+
+● Fetch(https://kalinga.ai/openenv-ai-hackathon-india-2026/)
+  ⎿  Received 193KB (200 OK)
+
+● Web Search("OpenEnv hackathon round 2 finale requirements training GRPO
+            submission themes judging 2026")
+  ⎿  Did 10 searches in 7s
+
+● Fetch(https://github.com/meta-pytorch/OpenEnv)
+  ⎿  Received 434.6KB (200 OK)
+
+● Web Search(""openenv" hackathon themes "world modeling" "long horizon" judging
+             criteria scoring environment 2026")
+  ⎿  Did 8 searches in 10s
+
+● Web Search("OpenEnv hackathon "round 2" OR "finale" training pipeline GRPO
+            unsloth TRL submission requirements environment")
+  ⎿  Did 6 searches in 11s
+
+● Fetch(https://pytorch.org/event/openenv-ai-hackathon/)
+  ⎿  Received 160.2KB (200 OK)
+
+● Fetch(https://huggingface.co/docs/trl/openenv)
+  ⎿  Received 27.2KB (200 OK)
+
+● Fetch(https://deepwiki.com/meta-pytorch/OpenEnv/12-examples-and-tutorials)
+  ⎿  Interrupted · What should Claude do instead?
+
+❯ If you are unable to read via url I'll provide you the raw text               
+  ⎿  UserPromptSubmit hook error                                       ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
+
+● Yes, please paste the raw text from the documents. I've gathered significant
+  context from public sources (Scaler, PyTorch.org, TRL docs, GitHub), but the
+  Google Docs are behind authentication and contain the critical finale-specific
+   details I need — especially the Themes & Judging Criteria document and the
+  Participant Help Guide.
+
+  Please paste the content from these documents:
+
+  1. Hackathon FAQs (1ysO...)
+  2. Themes & Judging Criteria (1Odz...)
+  3. OpenEnv Hackathon Resources (1m4E...)
+  4. Participant Help Guide (1lGO...)
+
+  Also, if you can share the key points from the YouTube video (Mega Lecture
+  91), that would fill the remaining gaps. I'll compile the full action plan
+  once I have the raw text.
+
+● Ran 1 stop hook (ctrl+o to expand)
+  ⎿  Stop hook error: Failed with non-blocking status code:
+  /usr/bin/bash: line 1: C:Program: command not found
+
+❯ I have added everything in @docs\                                             
+  ⎿  Listed directory docs\
+  ⎿  UserPromptSubmit hook error                                       ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
+
+● The documents are in the repo. Let me read all of them now.
+
+  Read 6 files (ctrl+o to expand)
+  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
+  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
+  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
+  ⎿  PostToolUseFailure:Read hook error                                ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: 
+
+● The Themes & Judging Criteria doc and the FAQ are large. Let me read them in
+  parts.
+
+  Searched for 1 pattern, read 2 files, listed 1 directory (ctrl+o to expand)  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:Bash hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+  ⎿  PostToolUseFailure:Bash hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+● Now reading the FAQ and the transcript (key sections).
+
+  Searched for 1 pattern, read 1 file (ctrl+o to expand)
+  ⎿  PostToolUseFailure:Read hook error  ⎿  Failed with non-blocking status code: /usr/bin/bash: line 1: C:Program: command not found
+
+● Now let me read the key portions of the transcript for hackathon-specific
+  guidance.
+
+  Searched for 1 pattern (ctrl+o to expand)
+
+● I now have all the critical information. Let me also verify the Themes
+  document is fully captured since I read it via bash.
+
+  Searched for 1 pattern, read 1 file (ctrl+o to expand)
+
+● I now have all the information needed. Here is the complete analysis:
+
+  ---
+  SECTION 1 — Final Round Requirements
+
+  Explicit organizer requirements:
+
+  1. Use OpenEnv (latest release) — build on top of the framework, don't
+  reinvent the wheel
+  2. A working training script using Unsloth or HF TRL, ideally as a Colab 
+  notebook so judges can re-run it
+  3. Evidence of actual training — at minimum, loss and reward plots from a real
+   run
+  4. A short writeup: mini-blog on HuggingFace OR < 2 minute video on YouTube,
+  OR a short slide deck
+  5. Environment hosted on Hugging Face Spaces — discoverable and runnable;
+  judges will pull the environment from the submitted URL
+  6. A README that motivates the problem, explains the environment, and shows
+  results — with links to HF Space, blog/video, and all supplementary materials
+  7. One submission per team — pick your best idea
+  8. No large video files in the HF Hub submission — use URL references instead
+  9. No changes/commits after the submission deadline — judges evaluate the
+  snapshot at deadline
+  10. It is NOT mandatory to keep the same problem statement as Round 1 — only
+  keep it if it aligns with the finale themes
+  11. Compute credits for HuggingFace will be given onsite on April 25-26 for
+  post-training
+
+  ---
+  SECTION 2 — Required Deliverables
+
+  #: 1
+  Deliverable: OpenEnv-compliant environment
+  Format: HF Space (live URL)
+  Status for MaterialForge: Partially done — deployed at
+    ArshPathan/material_forge_env but needs verification with latest OpenEnv
+  ────────────────────────────────────────
+  #: 2
+  Deliverable: Training script
+  Format: Colab notebook (Unsloth or TRL + GRPO)
+  Status for MaterialForge: NOT DONE — critical gap
+  ────────────────────────────────────────
+  #: 3
+  Deliverable: Training evidence
+  Format: Reward curves, loss plots, before/after comparison (.png/.jpg in repo
+  +
+    README)
+  Status for MaterialForge: NOT DONE — critical gap
+  ────────────────────────────────────────
+  #: 4
+  Deliverable: Mini-blog OR <2min video OR slide deck
+  Format: HF blog post / YouTube / Slides
+  Status for MaterialForge: NOT DONE
+  ────────────────────────────────────────
+  #: 5
+  Deliverable: README with results
+  Format: Markdown in repo root
+  Status for MaterialForge: Exists but needs major update with training results,
+
+    plots, and links
+  ────────────────────────────────────────
+  #: 6
+  Deliverable: openenv.yaml manifest
+  Format: YAML in repo root
+  Status for MaterialForge: Done
+  ────────────────────────────────────────
+  #: 7
+  Deliverable: Valid client/server separation
+  Format: Code architecture
+  Status for MaterialForge: Done
+  ────────────────────────────────────────
+  #: 8
+  Deliverable: 3-minute live pitch
+  Format: In-person presentation
+  Status for MaterialForge: NOT PREPARED
+
+  ---
+  SECTION 3 — Training Requirements
+
+  Is RL training expected?
+
+  YES — absolutely mandatory. The documents are explicit:
+
+  ▎ "The bar isn't 'training script exists.' The bar is 'training script runs 
+  ▎ against the environment, the agent learns, and you can show it.'"
+
+  Training evidence is worth 20% of the score directly, and the training
+  pipeline is another 10%.
+
+  Recommended stack:
+
+  - TRL (GRPOTrainer) — the primary training framework
+  - Unsloth — for memory efficiency (QLoRA/LoRA), especially on limited hardware
+  - GRPO algorithm — preferred over PPO for LLM RL; more memory-efficient, no
+  value model needed
+  - OpenEnv environment_factory pattern — wraps your environment as tool methods
+   the trainer discovers automatically
+
+  How training should work (from TRL docs):
+
+  1. Create an environment class with reset() and tool methods (e.g.,
+  place_atom(), remove_atom())
+  2. The GRPOTrainer discovers tool methods, generates completions, parses tool
+  calls, executes them, feeds results back
+  3. Reward comes from the environment after each episode
+  4. Model: Start with a capable instruct model (e.g., Qwen/Qwen3-0.6B or
+  Qwen/Qwen2.5-3B)
+
+  Expected evaluation metrics:
+
+  - Reward curves over training steps (must be plotted, labeled, saved as
+  images)
+  - Before/after behavior comparison (untrained baseline vs. trained agent)
+  - Loss plots
+  - Quantitative comparison: random agent vs. heuristic baseline vs. trained LLM
+   agent
+
+  Key warnings from organizers:
+
+  - Don't naively upcast 4-bit to 16-bit when saving LoRA — use proper merge
+  path
+  - Start with easy tasks, use curriculum learning
+  - Binary rewards (1.0 success / 0.0 fail) often work better than shaped
+  rewards for GRPO
+  - max_completion_length must be increased for long episodes (default is too
+  small)
+  - Test reward function before training — if a capable model can't beat random,
+   reward is broken
+
+  ---
+  SECTION 4 — Submission Process
+
+  Where submissions happen:
+
+  - HuggingFace Spaces — the environment must be live and accessible via URL
+  - Judges pull the environment from the submitted URL to evaluate it
+  - All supplementary materials (blog, video, slides, plots) must be linked from
+   the README
+
+  What files/links are needed:
+
+  1. HF Space URL of the running environment
+  2. Colab notebook with training script (linked from README)
+  3. Plot images (.png/.jpg) committed to the repo and embedded in README
+  4. Blog/video/slides link in README
+  5. openenv.yaml manifest in repo root
+  6. README.md with: problem motivation, how env works, results, and all links
+
+  Format requirements:
+
+  - Plots: labeled axes, .png or .jpg, embedded in README with one-line captions
+  - Video: < 2 minutes
+  - No large binary files in HF Hub
+  - README readable in 3-5 minutes
+
+  ---
+  SECTION 5 — Judging Criteria
+
+  Criterion: Environment Innovation
+  Weight: 40%
+  What judges evaluate: Is it novel, creative, genuinely challenging? Does it
+    teach an LLM something it can't currently do well? Underexplored domain?
+    Fresh angle? (Judges have seen too many chess/snake/grid-world clones)
+  ────────────────────────────────────────
+  Criterion: Storytelling & Presentation
+  Weight: 30%
+  What judges evaluate: Clear explanation of problem, environment, and what the
+    agent learned. Engaging demo, easy to follow for non-technical audience.
+    README answers: Problem → Environment → Results → Why it matters
+  ────────────────────────────────────────
+  Criterion: Showing Improvement in Rewards
+  Weight: 20%
+  What judges evaluate: Observable training progress: reward curves,
+  before/after
+    behavior, baseline comparison. Proves the agent actually learned something
+  ────────────────────────────────────────
+  Criterion: Reward & Training Pipeline
+  Weight: 10%
+  What judges evaluate: Coherent reward logic, pipeline produces meaningful
+    improvement in agent behavior
+
+  What makes a submission stand out:
+
+  - Ambitious, original problem — "Does this environment exist to teach an LLM
+  something it currently can't do well?"
+  - Rich reward signal — not just 0/1, uses OpenEnv's Rubric system, hard to
+  game
+  - Real end-to-end training — not just "script exists" but "script runs, agent
+  learns, you can show it"
+  - Readable plots on same axes for comparison
+  - A story, not an API doc — problem gap → environment → results → why it
+  matters
+
+  MaterialForge's theme alignment:
+
+  MaterialForge fits Theme #3.1 (World Modeling — Professional Tasks) and Theme 
+  #2 (Long-Horizon Planning). The scientific workflow (inverse molecular design)
+   is a genuinely novel domain for LLM RL training — this is a strong
+  differentiator.
+
+  ---
+  SECTION 6 — Immediate Action Plan
+
+  CRITICAL PATH (must-do, in priority order):
+
+  - 1. Build the TRL/GRPO Training Script (HIGHEST PRIORITY)
+    - Create a Colab notebook that uses GRPOTrainer with environment_factory
+    - Wrap MaterialForge as a TRL-compatible environment class with tool methods
+   (place_atom, remove_atom, replace_atom) instead of a generic step()
+    - Use Unsloth for memory efficiency
+    - Start model: Qwen/Qwen3-0.6B or Qwen/Qwen2.5-3B
+    - Connect to the MaterialForge HF Space via WebSocket
+    - This unlocks 30% of the score (20% improvement evidence + 10% pipeline)
+  - 2. Run Training & Capture Evidence
+    - Execute training with at least enough steps for curves to be meaningful
+    - Save reward curves and loss plots as .png with labeled axes
+    - Record baseline comparison: random agent score vs. heuristic vs. trained
+  LLM
+    - Use the onsite HF compute credits (April 25-26) for larger runs
+    - Before/after qualitative examples (grid screenshots showing trained agent
+  building better crystals)
+  - 3. Update README with Results
+    - Restructure to tell the story: Problem → Environment → Results → Why It
+  Matters
+    - Embed training plots with one-line captions
+    - Add links to: HF Space, Colab training notebook, video/blog
+    - Show baseline vs. trained comparison table
+    - Make it readable in 3-5 minutes
+  - 4. Create Mini-Blog or <2min Video
+    - Option A: HuggingFace blog post explaining the environment and training
+  results
+    - Option B: YouTube video (<2 min) showing the Discovery Lab dashboard +
+  training progress
+    - Option C: Short slide deck for the pitch
+    - Link it from README
+  - 5. Verify HF Space is Live & Current
+    - Ensure ArshPathan/material_forge_env is running and accessible
+    - Test that judges can hit /reset and /step endpoints
+    - Verify openenv.yaml is valid with latest OpenEnv release
+  - 6. Prepare 3-Minute Pitch
+    - Structure: Problem (inverse molecular design) → Environment (what agent
+  sees/does) → Demo (show dashboard) → Training Results (reward curves,
+  before/after) → Why it matters (scientific discovery)
+    - Practice delivering in 3 minutes
+    - Emphasize the novelty: scientific professional workflow, not another game
+  clone
+
+  NICE-TO-HAVE (if time permits):
+
+  - Add curriculum learning (easy → medium → hard scenarios during training)
+  - Multiple reward function ablations to show robustness
+  - Anti-reward-hacking analysis
+  - WandB integration for richer training logs
+
+  Time allocation suggestion (for the 48-hour finale):
+
+  ┌───────────┬───────┬─────────────────────────────────────────────────────┐
+  │   Phase   │ Hours │                      Activity                       │
+  ├───────────┼───────┼────��────────────────────────────────────────────────┤
+  │ Hours 0-4 │ 4h    │ Build TRL training script + environment wrapper in  │
+  │           │       │ Colab                                               │
+  ├───────────┼───────┼─────────────────────────────────────────────────────┤
+  │ Hours 4-8 │ 4h    │ Debug training loop, get first reward curves        │
+  ├───────────┼───────┼─────────────────────────────────────────────────────┤
+  │ Hours     │ 8h    │ Run training with onsite compute credits, iterate   │
+  │ 8-16      │       │ on rewards                                          │
+  ├───────────┼───────┼─────────────────────────────────────────────────────┤
+  │ Hours     │ 8h    │ Capture all evidence, generate plots, run baseline  │
+  │ 16-24     │       │ comparisons                                         │
+  ├───────────┼───────┼─────────────────────────────────────────────────────┤
+  │ Hours     │ 8h    │ Update README, write blog/create video, clean up    │
+  │ 24-32     │       │ repo                                                │
+  ├───────────┼───────┼─────────────────────────────────────────────────────┤
+  │ Hours     │ 8h    │ Polish demo, prepare pitch, verify HF Space         │
+  │ 32-40     │       │                                                     │
+  ├───────────┼───────┼─────────────────────────────────────────────────────┤
+  │ Hours     │ 8h    │ Final submission prep, rehearse pitch, buffer for   │
+  │ 40-48     │       │ issues                                              │
+  └───────────┴───────┴─────────────────────────────────────────────────────┘
+
+● User answered Claude's questions:
+  ⎿  · [Bonsai] What went well or could be improved on this task? Consider 
+     correctness, performance, and code quality. → Awsome performance
+
+● User answered Claude's questions:
+  ⎿  · [Bonsai] How did the model do on this task? → Great
+
+● Thank you for the feedback!
+
+● Ran 1 stop hook (ctrl+o to expand)
+  ⎿  Stop hook error: Failed with non-blocking status code:
+  /usr/bin/bash: line 1: C:Program: command not found
+
 ✻ Cogitated for 4m 42s

environment/config.py

@@ -14,34 +14,35 @@ VALID_COMPLETION_PHASES = {"crystalline", "polycrystalline"}
 PROPERTY_NAMES = ["hardness", "conductivity", "thermal_resistance", "elasticity"]
 
 # Atom definitions: defines species symbols, their resource costs, and physical contributions.
+# Each species is a SPECIALIST in one physical property to force chemical diversity.
 ATOM_TYPES = {
     "A": {
         "name": "metal",
         "cost": 8,
         "contributions": {
-            "hardness": 1.00,
-            "conductivity": 0.45,
-            "thermal_resistance": 0.35,
-            "elasticity": 0.10,
+            "hardness": 1.00,           # Master of Hardness
+            "conductivity": 0.15,
+            "thermal_resistance": 0.10,
+            "elasticity": 0.05,
         },
     },
     "B": {
         "name": "conductor",
         "cost": 6,
         "contributions": {
-            "hardness": 0.20,
-            "conductivity": 1.00,
-            "thermal_resistance": 0.15,
-            "elasticity": 0.30,
+            "hardness": 0.10,
+            "conductivity": 1.00,       # Master of Conductivity
+            "thermal_resistance": 0.10,
+            "elasticity": 0.20,
         },
     },
     "C": {
         "name": "ceramic",
         "cost": 4,
         "contributions": {
-            "hardness": 0.65,
+            "hardness": 0.30,
             "conductivity": 0.05,
-            "thermal_resistance": 1.00,
+            "thermal_resistance": 1.00, # Master of Thermal Resistance
             "elasticity": 0.05,
         },
     },
@@ -49,10 +50,10 @@ ATOM_TYPES = {
         "name": "polymer",
         "cost": 2,
         "contributions": {
-            "hardness": 0.10,
-            "conductivity": 0.10,
-            "thermal_resistance": 0.20,
-            "elasticity": 1.00,
+            "hardness": 0.05,
+            "conductivity": 0.15,
+            "thermal_resistance": 0.15,
+            "elasticity": 1.00,         # Master of Elasticity
         },
     },
 }

training/MaterialForge_GRPO_Training.ipynb

@@ -4,468 +4,104 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# MaterialForge: GRPO Training for Crystal Design\n",
+    "# MaterialForge: Reasoning-Based GRPO Training (100 Episode Sprint)\n",
     "\n",
-    "**Meta PyTorch OpenEnv Hackathon x Scaler School of Technology \u2014 Grand Finale**\n",
+    "**Meta PyTorch OpenEnv Hackathon \u2014 Grand Finale**\n",
     "\n",
-    "This notebook trains an LLM to autonomously design crystalline structures using reinforcement learning (GRPO) with the MaterialForge environment.\n",
-    "\n",
-    "| Component | Details |\n",
-    "|-----------|----------|\n",
-    "| **Environment** | MaterialForge \u2014 8x8 atomic lattice, 4 atom species, 4 physical properties |\n",
-    "| **Algorithm** | GRPO (Group Relative Policy Optimization) via HuggingFace TRL |\n",
-    "| **Efficiency** | Unsloth for 4-bit QLoRA training |\n",
-    "| **Framework** | OpenEnv-compliant environment |\n",
-    "| **HF Space** | [ArshPathan/material_forge_env](https://huggingface.co/spaces/ArshPathan/material_forge_env) |"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 1. Installation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "metadata": {
-    "jupyter": {
-     "is_executing": true
-    }
-   },
-   "source": "%%capture\n!pip install -Uq unsloth\n!pip install -Uq trl>=0.18.0 transformers datasets accelerate peft\n!pip install -Uq \"openenv-core[core]>=0.2.3\" aiofiles openai python-dotenv\n!pip install -Uq matplotlib jmespath",
-   "outputs": [],
-   "execution_count": null
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 2. Clone Environment & Setup Paths\n",
-    "\n",
-    "We clone the MaterialForge repository to access the environment code directly.\n",
-    "The same environment is hosted live at: https://huggingface.co/spaces/ArshPathan/material_forge_env"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "REPO_URL = \"https://huggingface.co/spaces/ArshPathan/material_forge_env\"\n",
-    "REPO_DIR = \"MaterialForge\"\n",
-    "\n",
-    "if not os.path.isdir(REPO_DIR):\n",
-    "    !git clone {REPO_URL} {REPO_DIR}\n",
-    "else:\n",
-    "    print(f\"{REPO_DIR}/ already exists, skipping clone.\")\n",
-    "\n",
-    "import sys\n",
-    "sys.path.insert(0, REPO_DIR)\n",
-    "\n",
-    "print(\"Environment source ready.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 3. Verify Environment Works\n",
-    "\n",
-    "Quick sanity check: reset the environment, take one action, confirm reward is computed."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from environment.rubrics import HeuristicRewardRubric\n",
-    "from server.material_forge_env_environment import MaterialForgeEnvironment\n",
-    "from models import MaterialForgeAction\n",
-    "\n",
-    "env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n",
-    "obs = env.reset(seed=42, difficulty=\"easy\")\n",
-    "print(f\"Target properties: {obs.target}\")\n",
-    "print(f\"Cost budget: {obs.cost_budget}\")\n",
-    "print(f\"Max steps: {obs.max_steps}\")\n",
-    "\n",
-    "action = MaterialForgeAction(action_type=\"place\", row=3, col=3, atom=\"A\")\n",
-    "obs = env.step(action)\n",
-    "print(f\"\\nAfter placing Metal at (3,3):\")\n",
-    "print(f\"  Reward: {obs.reward}\")\n",
-    "print(f\"  Phase: {obs.phase}\")\n",
-    "print(f\"  Properties: {obs.current_properties}\")\n",
-    "print(\"\\nEnvironment OK!\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 4. TRL Environment Wrapper\n",
-    "\n",
-    "TRL's `GRPOTrainer` with `environment_factory` requires a wrapper class that exposes **named tool methods** (not a generic `step()`). The trainer discovers these methods and presents them to the LLM as callable tools.\n",
-    "\n",
-    "Our wrapper exposes three tools:\n",
-    "- `place_atom(row, col, atom)` \u2014 place a new atom on an empty cell\n",
-    "- `remove_atom(row, col)` \u2014 remove an existing atom\n",
-    "- `replace_atom(row, col, atom)` \u2014 swap an atom for a different species\n",
-    "\n",
-    "Each tool internally calls the OpenEnv environment's standard `step()` method."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": "import random\nfrom typing import Optional\n\nfrom environment.config import GRID_SIZE, PROPERTY_NAMES, ATOM_TYPES\nfrom environment.rubrics import HeuristicRewardRubric\nfrom models import MaterialForgeAction, MaterialForgeObservation\nfrom server.material_forge_env_environment import MaterialForgeEnvironment\n\n\nDIFFICULTIES = [\"easy\", \"medium\", \"hard\"]\n\n\nclass MaterialForgeTRLEnv:\n    \"\"\"TRL-compatible wrapper for MaterialForge crystal design.\"\"\"\n\n    def __init__(self):\n        self.env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n        self.reward = 0.0\n        self.done = False\n        self._obs = None\n        self._best_reward = 0.0\n        self._reward_history = []\n        self._invalid_actions = 0\n        self._total_actions = 0\n\n    def reset(self, **kwargs) -> str:\n        seed = random.randint(0, 99999)\n        difficulty = random.choice(DIFFICULTIES)\n        obs = self.env.reset(seed=seed, difficulty=difficulty)\n        self._obs = obs\n        self.reward = 0.0\n        self.done = False\n        self._best_reward = 0.0\n        self._reward_history = []\n        self._invalid_actions = 0\n        self._total_actions = 0\n        return self._format_observation(obs)\n\n    def place_atom(self, row: int, col: int, atom: str) -> str:\n        \"\"\"Place an atom on an empty cell of the crystal lattice.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n            atom: Atom species - 'A' (Metal, hardness, cost 8), 'B' (Conductor, conductivity, cost 6), 'C' (Ceramic, thermal resistance, cost 4), or 'P' (Polymer, elasticity, cost 2)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"place\", row, col, atom)\n\n    def remove_atom(self, row: int, col: int) -> str:\n        \"\"\"Remove an atom from the crystal lattice to free up budget.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"remove\", row, col, None)\n\n    def replace_atom(self, row: int, col: int, atom: str) -> str:\n        \"\"\"Replace an existing atom with a different species to adjust material properties.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n            atom: New atom species - 'A' (Metal), 'B' (Conductor), 'C' (Ceramic), or 'P' (Polymer)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"replace\", row, col, atom)\n\n    def _do_step(self, action_type: str, row: int, col: int, atom: Optional[str]) -> str:\n        if self.done:\n            raise ValueError(\"Episode is over. The crystal design session has ended.\")\n\n        self._total_actions += 1\n        row = max(0, min(int(row), GRID_SIZE - 1))\n        col = max(0, min(int(col), GRID_SIZE - 1))\n\n        if atom is not None:\n            atom = str(atom).upper().strip()\n            if atom not in ATOM_TYPES:\n                atom = \"A\"\n\n        # Pre-check: will this action actually change the grid?\n        grid = self._obs.grid if self._obs else None\n        if grid:\n            cell = grid[row][col]\n            if action_type == \"place\" and cell != \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot place at ({row},{col}) \u2014 cell is already \"\n                    f\"occupied by '{cell}'. Use replace_atom to change it, or pick an empty cell (marked '.').\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"remove\" and cell == \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot remove from ({row},{col}) \u2014 cell is empty. \"\n                    f\"Pick a cell that contains an atom.\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"replace\" and cell == \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot replace at ({row},{col}) \u2014 cell is empty. \"\n                    f\"Use place_atom instead.\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"replace\" and cell == atom:\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cell ({row},{col}) already contains '{atom}'. \"\n                    f\"Pick a different atom species to replace it with.\\n\"\n                    + self._format_observation(self._obs)\n                )\n\n        action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)\n        obs = self.env.step(action)\n        self._obs = obs\n        step_reward = obs.reward if obs.reward is not None else 0.0\n        self._reward_history.append(step_reward)\n        self._best_reward = max(self._best_reward, step_reward)\n        self.reward = self._best_reward\n        self.done = obs.done\n        return self._format_observation(obs)\n\n    def _format_observation(self, obs) -> str:\n        grid_lines = []\n        for r, row in enumerate(obs.grid):\n            grid_lines.append(f\"  {r}: {' '.join(cell if cell != '.' else '.' for cell in row)}\")\n        grid_str = \"\\n\".join(grid_lines)\n\n        gaps = {}\n        for prop in PROPERTY_NAMES:\n            t = obs.target.get(prop, 0.0)\n            c = obs.current_properties.get(prop, 0.0)\n            gaps[prop] = round(t - c, 1)\n\n        props_str = \", \".join(\n            f\"{p}={obs.current_properties.get(p, 0):.1f}/{obs.target.get(p, 0):.1f}(gap:{gaps[p]:+.1f})\"\n            for p in PROPERTY_NAMES\n        )\n\n        sb = obs.score_breakdown or {}\n        stability = sb.get(\"structural_stability\", 0.0)\n        order = sb.get(\"lattice_order_index\", 0.0)\n\n        return (\n            f\"Step {obs.step_number}/{obs.max_steps} | \"\n            f\"Cost: {obs.total_cost:.0f}/{obs.cost_budget:.0f} | \"\n            f\"Phase: {obs.phase} | \"\n            f\"Reward: {obs.reward:.4f}\\n\"\n            f\"Properties: {props_str}\\n\"\n            f\"Stability: {stability:.3f} | Order: {order:.3f}\\n\"\n            f\"Grid (row: cols 0-7):\\n{grid_str}\"\n        )\n\n\nprint(\"MaterialForgeTRLEnv defined.\")"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Test the wrapper"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "test_env = MaterialForgeTRLEnv()\n",
-    "init_obs = test_env.reset()\n",
-    "print(\"=== Initial State ===\")\n",
-    "print(init_obs)\n",
-    "\n",
-    "print(\"\\n=== After place_atom(3, 3, 'A') ===\")\n",
-    "result = test_env.place_atom(3, 3, \"A\")\n",
-    "print(result)\n",
-    "print(f\"\\nReward stored: {test_env.reward:.4f}\")\n",
-    "print(f\"Done: {test_env.done}\")\n",
-    "\n",
-    "print(\"\\n=== After place_atom(3, 4, 'B') ===\")\n",
-    "result = test_env.place_atom(3, 4, \"B\")\n",
-    "print(result)\n",
-    "print(f\"Reward: {test_env.reward:.4f}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 5. Reward Function"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def reward_func(environments, **kwargs) -> list[float]:\n    \"\"\"Reward = best_reward + spatial_bonus + phase_bonus - invalid_penalty.\"\"\"\n    rewards = []\n    for env in environments:\n        best = env._best_reward\n        total = max(env._total_actions, 1)\n\n        # Penalty for invalid actions\n        invalid_ratio = env._invalid_actions / total\n        invalid_penalty = 0.3 * invalid_ratio\n\n        # Spatial diversity bonus: reward using multiple rows and columns\n        spatial_bonus = 0.0\n        if env._obs is not None:\n            grid = env._obs.grid\n            rows_used = set()\n            cols_used = set()\n            for r, row in enumerate(grid):\n                for c, cell in enumerate(row):\n                    if cell != \".\":\n                        rows_used.add(r)\n                        cols_used.add(c)\n            # Bonus for spreading across rows AND columns (max 0.10)\n            row_spread = min(len(rows_used) / 4.0, 1.0)  # 4+ rows = full bonus\n            col_spread = min(len(cols_used) / 4.0, 1.0)  # 4+ cols = full bonus\n            spatial_bonus = 0.10 * (row_spread * col_spread)\n\n        # Phase bonus\n        phase_bonus = 0.0\n        if env._obs is not None:\n            phase = env._obs.phase\n            if phase == \"crystalline\":\n                phase_bonus = 0.15\n            elif phase == \"polycrystalline\":\n                phase_bonus = 0.05\n\n        final = best + spatial_bonus + phase_bonus - invalid_penalty\n        rewards.append(max(min(final, 1.0), 0.0))\n    return rewards\n\nprint(\"Reward function: best_reward + spatial_bonus + phase_bonus - invalid_penalty\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 6. Training Dataset\n",
-    "\n",
-    "We create a prompt dataset where each entry instructs the LLM to design a crystal.\n",
-    "The environment generates random targets on each `reset()`, so every episode is unique."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from datasets import Dataset\n\nSYSTEM_PROMPT = \"\"\"You are an expert materials scientist designing crystal structures on an 8x8 atomic lattice.\n\nYour goal: place atoms to match target physical properties while staying within budget.\n\nAvailable atom species:\n- A (Metal): Primary contribution to hardness, cost 8 per atom\n- B (Conductor): Primary contribution to conductivity, cost 6 per atom\n- C (Ceramic): Primary contribution to thermal resistance, cost 4 per atom\n- P (Polymer): Primary contribution to elasticity, cost 2 per atom\n\nStrategy:\n- START from the center of the grid (around row 3-4, col 3-4) and expand outward in all directions\n- Build 2D clusters, NOT single rows or columns \u2014 atoms need neighbors in multiple directions for bonding\n- Group same-type atoms together in 2x2 or 3x3 blocks for maximum bonding bonuses\n- After each placement, READ the property gaps and choose the atom type that closes the LARGEST gap\n- Aim for crystalline phase: fill a compact rectangular region, not scattered atoms\n- Stay within the cost budget \u2014 use cheaper atoms (P=2, C=4) when budget is tight\n- Use replace_atom to fix overshooting properties, remove_atom if over budget\n\nUse the tools place_atom, remove_atom, and replace_atom to build your crystal.\"\"\"\n\nUSER_PROMPT = \"Design a crystal lattice that matches the target material properties. Start placing atoms from the center of the grid (row 3-4, col 3-4) and expand outward in all directions. Analyze the property gaps after each action and choose your next atom strategically.\"\n\nNUM_EPISODES = 200\n\ndataset = Dataset.from_dict({\n    \"prompt\": [\n        [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\", \"content\": USER_PROMPT},\n        ]\n    ] * NUM_EPISODES\n})\n\nprint(f\"Dataset created: {len(dataset)} episodes\")\nprint(f\"Sample prompt roles: {[m['role'] for m in dataset[0]['prompt']]}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 7. Load Model with Unsloth\n",
-    "\n",
-    "We use Unsloth for memory-efficient 4-bit QLoRA fine-tuning. Starting with `Qwen/Qwen3-0.6B` which fits on a free Colab T4 GPU. For better results with more compute, switch to `Qwen/Qwen3-1.7B` or `Qwen/Qwen2.5-3B-Instruct`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": "import torch\nfrom unsloth import FastLanguageModel\n\nMODEL_NAME = \"Qwen/Qwen3-0.6B\"  # Change to Qwen3-1.7B or Qwen2.5-3B-Instruct for better results\nMAX_SEQ_LENGTH = 4096\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    MODEL_NAME,\n    max_seq_length=MAX_SEQ_LENGTH,\n    load_in_4bit=True,\n    dtype=None,\n)\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=16,\n    lora_alpha=16,\n    lora_dropout=0,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n    use_gradient_checkpointing=\"unsloth\",\n)\n\n# --- Fix Unsloth + TRL compatibility for Qwen3 tool calling ---\n# Unsloth modifies the chat template, breaking TRL's string-comparison checks.\n# We load TRL's own Qwen3 training template (prefix-preserving, has {% generation %}\n# markers) and response schema directly.\nfrom pathlib import Path\nimport trl.chat_template_utils as _ctu\n\n# 1) Set the response schema so TRL can parse <tool_call> blocks\ntokenizer.response_schema = {\n    \"x-regex\": r\"^(?:<think>\\n?(?:(?P<reasoning_content>.*?\\S.*?)\\n?|[\\s]*)</think>\\s*)?(?P<content>.*?)(?:\\n(?=<tool_call>))?(?=(?:<tool_call>|<\\|im_end\\|>|$))(?P<tool_calls>(?:<tool_call>.+?</tool_call>\\s*)+)?\\s*(?:<\\|im_end\\|>|$)\",\n    \"type\": \"object\",\n    \"properties\": {\n        \"role\": {\"const\": \"assistant\"},\n        \"content\": {\"type\": \"string\"},\n        \"reasoning_content\": {\"type\": \"string\"},\n        \"tool_calls\": {\n            \"type\": \"array\",\n            \"x-regex-iterator\": r\"<tool_call>\\s*(.+?)\\s*</tool_call>\",\n            \"items\": {\n                \"x-parser\": \"json\",\n                \"x-parser-args\": {\"transform\": \"{type: 'function', function: @}\"},\n                \"type\": \"object\",\n                \"properties\": {\n                    \"type\": {\"const\": \"function\"},\n                    \"function\": {\n                        \"type\": \"object\",\n                        \"properties\": {\n                            \"name\": {\"type\": \"string\"},\n                            \"arguments\": {\"type\": \"object\", \"additionalProperties\": {}},\n                        },\n                    },\n                },\n            },\n        },\n    },\n}\n\n# 2) Replace Unsloth's modified template with TRL's training-safe Qwen3 template\n_templates_dir = Path(_ctu.__file__).parent / \"chat_templates\"\ntokenizer.chat_template = (_templates_dir / \"qwen3_training.jinja\").read_text()\n\nprint(f\"Model loaded: {MODEL_NAME}\")\nprint(f\"Trainable params: {model.print_trainable_parameters()}\")\nprint(f\"Chat template: TRL qwen3_training.jinja (prefix-preserving)\")\nprint(f\"Response schema: Qwen3 tool-call parser\")"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 8. GRPO Training\n",
-    "\n",
-    "The `GRPOTrainer` with `environment_factory` handles the full multi-turn loop:\n",
-    "1. Generates completions from the LLM\n",
-    "2. Parses tool calls (`place_atom`, `remove_atom`, `replace_atom`)\n",
-    "3. Executes them against the MaterialForge environment\n",
-    "4. Feeds observations back to the model\n",
-    "5. Collects rewards and updates the policy via GRPO"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from trl import GRPOConfig, GRPOTrainer\n\nOUTPUT_DIR = \"materialforge-grpo-checkpoints\"\n\ntraining_args = GRPOConfig(\n    output_dir=OUTPUT_DIR,\n\n    # Generation\n    num_generations=4,\n    max_completion_length=2048,\n    temperature=1.0,\n\n    # Optimization\n    learning_rate=5e-5,\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=4,\n    num_train_epochs=1,\n\n    # Logging\n    logging_steps=1,\n    log_completions=True,\n    num_completions_to_print=1,\n\n    # Tool calling\n    chat_template_kwargs={\"enable_thinking\": False},\n\n    # Saving\n    save_steps=50,\n    save_total_limit=2,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_func,\n    train_dataset=dataset,\n    args=training_args,\n    environment_factory=MaterialForgeTRLEnv,\n)\n\nprint(\"Trainer initialized. Starting GRPO training...\")\nprint(f\"  Episodes: {NUM_EPISODES}\")\nprint(f\"  Generations per prompt: {training_args.num_generations}\")\nprint(f\"  Max completion length: {training_args.max_completion_length}\")\nprint(f\"  Temperature: {training_args.temperature}\")\nprint(f\"  Learning rate: {training_args.learning_rate}\")\nprint(f\"  Gradient accumulation: {training_args.gradient_accumulation_steps}\")\nprint(f\"  Total steps: {NUM_EPISODES // (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps)}\")"
+    "This notebook trains a **DeepSeek-R1 Reasoning Model** to design specialized crystalline lattices using specialist atomic species."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
-    "train_result = trainer.train()\n",
-    "print(f\"\\nTraining complete!\")\n",
-    "print(f\"  Total steps: {trainer.state.global_step}\")\n",
-    "print(f\"  Final loss: {train_result.training_loss:.4f}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 9. Save Trained Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
+    "%%capture\n",
+    "!pip install -Uq unsloth trl>=0.18.0 transformers datasets accelerate peft openenv-core[core]>=0.2.3 matplotlib jmespath"
+   ],
    "outputs": [],
-   "source": [
-    "SAVE_DIR = \"materialforge-grpo-trained\"\n",
-    "trainer.save_model(SAVE_DIR)\n",
-    "tokenizer.save_pretrained(SAVE_DIR)\n",
-    "print(f\"Model saved to {SAVE_DIR}/\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 10. Training Reward Curves\n",
-    "\n",
-    "Extract and plot the reward history from training logs. These plots are a **required deliverable** \u2014 they must be saved as images and embedded in the README."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import json\n",
-    "\n",
-    "os.makedirs(\"plots\", exist_ok=True)\n",
-    "\n",
-    "log_history = trainer.state.log_history\n",
-    "\n",
-    "steps = []\n",
-    "rewards = []\n",
-    "losses = []\n",
-    "\n",
-    "for entry in log_history:\n",
-    "    if \"reward\" in entry:\n",
-    "        steps.append(entry.get(\"step\", len(steps)))\n",
-    "        rewards.append(entry[\"reward\"])\n",
-    "    if \"loss\" in entry:\n",
-    "        if len(losses) < len(steps):\n",
-    "            losses.append(entry[\"loss\"])\n",
-    "\n",
-    "# --- Plot 1: Reward over training ---\n",
-    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
-    "ax.plot(steps[:len(rewards)], rewards, color=\"#4285f4\", linewidth=2, label=\"GRPO Reward\")\n",
-    "if len(rewards) > 5:\n",
-    "    window = max(3, len(rewards) // 10)\n",
-    "    smoothed = [sum(rewards[max(0,i-window):i+1])/len(rewards[max(0,i-window):i+1]) for i in range(len(rewards))]\n",
-    "    ax.plot(steps[:len(smoothed)], smoothed, color=\"#ea4335\", linewidth=2, linestyle=\"--\", label=f\"Smoothed (window={window})\")\n",
-    "ax.set_xlabel(\"Training Step\", fontsize=12)\n",
-    "ax.set_ylabel(\"Reward\", fontsize=12)\n",
-    "ax.set_title(\"MaterialForge GRPO Training \u2014 Reward Curve\", fontsize=14, fontweight=\"bold\")\n",
-    "ax.legend()\n",
-    "ax.grid(True, alpha=0.3)\n",
-    "fig.tight_layout()\n",
-    "fig.savefig(\"plots/reward_curve.png\", dpi=150)\n",
-    "plt.show()\n",
-    "print(\"Saved: plots/reward_curve.png\")\n",
-    "\n",
-    "# --- Plot 2: Loss over training ---\n",
-    "if losses:\n",
-    "    fig2, ax2 = plt.subplots(1, 1, figsize=(10, 5))\n",
-    "    ax2.plot(range(len(losses)), losses, color=\"#34a853\", linewidth=2)\n",
-    "    ax2.set_xlabel(\"Training Step\", fontsize=12)\n",
-    "    ax2.set_ylabel(\"Loss\", fontsize=12)\n",
-    "    ax2.set_title(\"MaterialForge GRPO Training \u2014 Loss Curve\", fontsize=14, fontweight=\"bold\")\n",
-    "    ax2.grid(True, alpha=0.3)\n",
-    "    fig2.tight_layout()\n",
-    "    fig2.savefig(\"plots/loss_curve.png\", dpi=150)\n",
-    "    plt.show()\n",
-    "    print(\"Saved: plots/loss_curve.png\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 11. Baseline Comparison\n",
-    "\n",
-    "Compare the trained agent against a random baseline to demonstrate improvement."
-   ]
+   "execution_count": null
   },
   {
    "cell_type": "code",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
-    "import random as rng\n",
-    "\n",
-    "def run_random_baseline(n_episodes=20, max_steps=30):\n",
-    "    \"\"\"Run a random agent as baseline.\"\"\"\n",
-    "    env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n",
-    "    episode_rewards = []\n",
-    "    episode_best = []\n",
-    "\n",
-    "    for ep in range(n_episodes):\n",
-    "        obs = env.reset(seed=ep * 7, difficulty=\"easy\")\n",
-    "        rewards = []\n",
-    "        for step in range(max_steps):\n",
-    "            action_type = rng.choice([\"place\", \"place\", \"place\", \"replace\", \"remove\"])\n",
-    "            row = rng.randint(0, 7)\n",
-    "            col = rng.randint(0, 7)\n",
-    "            atom = rng.choice([\"A\", \"B\", \"C\", \"P\"]) if action_type != \"remove\" else None\n",
-    "            action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)\n",
-    "            obs = env.step(action)\n",
-    "            rewards.append(obs.reward if obs.reward else 0.0)\n",
-    "            if obs.done:\n",
-    "                break\n",
-    "        episode_rewards.append(sum(rewards) / len(rewards) if rewards else 0.0)\n",
-    "        episode_best.append(max(rewards) if rewards else 0.0)\n",
-    "\n",
-    "    return episode_rewards, episode_best\n",
-    "\n",
-    "print(\"Running random baseline (20 episodes)...\")\n",
-    "random_avg_rewards, random_best_rewards = run_random_baseline(20)\n",
-    "print(f\"Random Agent \u2014 Mean reward: {sum(random_avg_rewards)/len(random_avg_rewards):.4f}\")\n",
-    "print(f\"Random Agent \u2014 Mean best:   {sum(random_best_rewards)/len(random_best_rewards):.4f}\")"
+    "from training.train_env_wrapper import MaterialForgeTRLEnv\n",
+    "from datasets import Dataset\n",
+    "\n",
+    "def reward_func(environments, completions, **kwargs) -> list[float]:\n",
+    "    rewards = []\n",
+    "    for env, completion in zip(environments, completions):\n",
+    "        # 1. Base Scientific Reward\n",
+    "        score = env.reward \n",
+    "        \n",
+    "        # 2. Reasoning Bonus (Innovation)\n",
+    "        thought_bonus = 0.0\n",
+    "        if \"<think>\" in completion and \"</think>\" in completion:\n",
+    "            thought = completion.split(\"</think>\")[0].split(\"<think>\")[1]\n",
+    "            if len(thought.strip()) > 100:\n",
+    "                thought_bonus = 0.05\n",
+    "\n",
+    "        rewards.append(max(min(score + thought_bonus, 1.0), 0.0))\n",
+    "    return rewards\n",
+    "\n",
+    "SYSTEM_PROMPT = \"\"\"You are a Material Scientist. You MUST reason inside <think> tags before every action.\n",
+    "\n",
+    "Strategic Rules:\n",
+    "1. Atoms are SPECIALISTS. Metal (A) is for Hardness, Conductor (B) for Conductivity, etc. \n",
+    "2. You MUST combine different species to hit multi-property targets.\n",
+    "3. Build compact 2D clusters starting from the center (3,3) to maximize stability.\n",
+    "4. Pure crystalline blocks are highly stable; avoid 1D lines.\"\"\"\n",
+    "\n",
+    "# Requesting 100 episodes for a fast finale sprint\n",
+    "NUM_EPISODES = 100\n",
+    "dataset = Dataset.from_dict({\"prompt\": [[\n",
+    "    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "    {\"role\": \"user\", \"content\": \"Design a crystal lattice. Think step-by-step and mix specialist atoms.\"}\n",
+    "]] * NUM_EPISODES})\n",
+    "print(f\"Dataset ready for {NUM_EPISODES} reasoning episodes.\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
-    "# Comparison bar chart\n",
-    "trained_mean_reward = sum(rewards) / len(rewards) if rewards else 0.0\n",
-    "random_mean_reward = sum(random_avg_rewards) / len(random_avg_rewards)\n",
+    "import torch\n",
+    "from unsloth import FastLanguageModel\n",
     "\n",
-    "fig3, ax3 = plt.subplots(figsize=(8, 5))\n",
-    "agents = [\"Random\\nBaseline\", \"GRPO-Trained\\nAgent\"]\n",
-    "means = [random_mean_reward, trained_mean_reward]\n",
-    "colors = [\"#ea4335\", \"#4285f4\"]\n",
-    "bars = ax3.bar(agents, means, color=colors, width=0.5, edgecolor=\"white\", linewidth=1.5)\n",
-    "\n",
-    "for bar, val in zip(bars, means):\n",
-    "    ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,\n",
-    "             f\"{val:.4f}\", ha=\"center\", va=\"bottom\", fontweight=\"bold\", fontsize=13)\n",
-    "\n",
-    "ax3.set_ylabel(\"Mean Reward\", fontsize=12)\n",
-    "ax3.set_title(\"MaterialForge \u2014 Random Baseline vs GRPO-Trained Agent\", fontsize=14, fontweight=\"bold\")\n",
-    "ax3.set_ylim(0, max(means) * 1.3 + 0.05)\n",
-    "ax3.grid(axis=\"y\", alpha=0.3)\n",
-    "fig3.tight_layout()\n",
-    "fig3.savefig(\"plots/baseline_comparison.png\", dpi=150)\n",
-    "plt.show()\n",
-    "print(\"Saved: plots/baseline_comparison.png\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 12. Trained Agent Inference Demo\n",
+    "MODEL_NAME = \"unsloth/DeepSeek-R1-Distill-Qwen-1.5B\"\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(MODEL_NAME, max_seq_length=2048, load_in_4bit=True)\n",
     "\n",
-    "Run the trained model on a fresh episode to see its crystal design strategy."
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model, r=16, lora_alpha=16, lora_dropout=0,\n",
+    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    use_gradient_checkpointing=\"unsloth\"\n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
-    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "from trl import GRPOConfig, GRPOTrainer\n",
     "\n",
-    "print(\"Loading trained model for inference...\")\n",
-    "inf_model, inf_tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    SAVE_DIR,\n",
-    "    max_seq_length=MAX_SEQ_LENGTH,\n",
-    "    load_in_4bit=True,\n",
+    "training_args = GRPOConfig(\n",
+    "    output_dir=\"materialforge-r1-checkpoints\",\n",
+    "    num_generations=2,\n",
+    "    max_completion_length=1024,\n",
+    "    learning_rate=5e-5,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=4, # Reduced for faster updates in a 100-episode run\n",
+    "    num_train_epochs=1,\n",
+    "    chat_template_kwargs={\"enable_thinking\": True},\n",
+    "    logging_steps=1,\n",
     ")\n",
-    "FastLanguageModel.for_inference(inf_model)\n",
     "\n",
-    "demo_env = MaterialForgeTRLEnv()\n",
-    "obs_text = demo_env.reset()\n",
-    "\n",
-    "messages = [\n",
-    "    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "    {\"role\": \"user\", \"content\": USER_PROMPT},\n",
-    "    {\"role\": \"assistant\", \"content\": f\"Let me analyze the initial state:\\n{obs_text}\\n\\nI'll start placing atoms strategically.\"},\n",
-    "]\n",
-    "\n",
-    "print(\"=\" * 60)\n",
-    "print(\"TRAINED AGENT DEMO\")\n",
-    "print(\"=\" * 60)\n",
-    "print(f\"\\nInitial observation:\\n{obs_text}\")\n",
-    "print(f\"\\nAgent designing crystal...\")\n",
-    "print(f\"(Full multi-turn inference requires TRL's tool-call loop.\")\n",
-    "print(f\"See the training logs above for actual agent behavior.)\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 13. Summary\n",
-    "\n",
-    "| Metric | Random Baseline | GRPO-Trained Agent |\n",
-    "|--------|----------------|--------------------|\n",
-    "| Mean Reward | *see above* | *see above* |\n",
-    "| Phase Achieved | Amorphous | *see training logs* |\n",
-    "\n",
-    "### Deliverables produced by this notebook:\n",
-    "- `plots/reward_curve.png` \u2014 Reward over training steps\n",
-    "- `plots/loss_curve.png` \u2014 Training loss curve\n",
-    "- `plots/baseline_comparison.png` \u2014 Random vs trained comparison\n",
-    "- `materialforge-grpo-trained/` \u2014 Saved model checkpoint\n",
+    "trainer = GRPOTrainer(\n",
+    "    model=model, processing_class=tokenizer, reward_funcs=reward_func,\n",
+    "    train_dataset=dataset, args=training_args, environment_factory=MaterialForgeTRLEnv\n",
+    ")\n",
     "\n",
-    "### Links\n",
-    "- **Live Environment**: [ArshPathan/material_forge_env](https://huggingface.co/spaces/ArshPathan/material_forge_env)\n",
-    "- **Repository**: [MaterialForge GitHub](https://github.com/Arsh-Pathan/MaterialForge)"
+    "trainer.train()"
    ]
   }
  ],
@@ -473,8 +109,7 @@
   "accelerator": "GPU",
   "colab": {
    "gpuType": "T4",
-   "provenance": [],
-   "name": "MaterialForge_GRPO_Training.ipynb"
+   "name": "MaterialForge_Reasoning_Training.ipynb"
   },
   "kernelspec": {
    "display_name": "Python 3",
@@ -487,4 +122,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}

training/Traning Runs/Run - II/MaterialForge_GRPO_Training.ipynb

@@ -0,0 +1,490 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# MaterialForge: GRPO Training for Crystal Design\n",
+    "\n",
+    "**Meta PyTorch OpenEnv Hackathon x Scaler School of Technology \u2014 Grand Finale**\n",
+    "\n",
+    "This notebook trains an LLM to autonomously design crystalline structures using reinforcement learning (GRPO) with the MaterialForge environment.\n",
+    "\n",
+    "| Component | Details |\n",
+    "|-----------|----------|\n",
+    "| **Environment** | MaterialForge \u2014 8x8 atomic lattice, 4 atom species, 4 physical properties |\n",
+    "| **Algorithm** | GRPO (Group Relative Policy Optimization) via HuggingFace TRL |\n",
+    "| **Efficiency** | Unsloth for 4-bit QLoRA training |\n",
+    "| **Framework** | OpenEnv-compliant environment |\n",
+    "| **HF Space** | [ArshPathan/material_forge_env](https://huggingface.co/spaces/ArshPathan/material_forge_env) |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Installation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "jupyter": {
+     "is_executing": true
+    }
+   },
+   "source": "%%capture\n!pip install -Uq unsloth\n!pip install -Uq trl>=0.18.0 transformers datasets accelerate peft\n!pip install -Uq \"openenv-core[core]>=0.2.3\" aiofiles openai python-dotenv\n!pip install -Uq matplotlib jmespath",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Clone Environment & Setup Paths\n",
+    "\n",
+    "We clone the MaterialForge repository to access the environment code directly.\n",
+    "The same environment is hosted live at: https://huggingface.co/spaces/ArshPathan/material_forge_env"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "REPO_URL = \"https://huggingface.co/spaces/ArshPathan/material_forge_env\"\n",
+    "REPO_DIR = \"MaterialForge\"\n",
+    "\n",
+    "if not os.path.isdir(REPO_DIR):\n",
+    "    !git clone {REPO_URL} {REPO_DIR}\n",
+    "else:\n",
+    "    print(f\"{REPO_DIR}/ already exists, skipping clone.\")\n",
+    "\n",
+    "import sys\n",
+    "sys.path.insert(0, REPO_DIR)\n",
+    "\n",
+    "print(\"Environment source ready.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Verify Environment Works\n",
+    "\n",
+    "Quick sanity check: reset the environment, take one action, confirm reward is computed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from environment.rubrics import HeuristicRewardRubric\n",
+    "from server.material_forge_env_environment import MaterialForgeEnvironment\n",
+    "from models import MaterialForgeAction\n",
+    "\n",
+    "env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n",
+    "obs = env.reset(seed=42, difficulty=\"easy\")\n",
+    "print(f\"Target properties: {obs.target}\")\n",
+    "print(f\"Cost budget: {obs.cost_budget}\")\n",
+    "print(f\"Max steps: {obs.max_steps}\")\n",
+    "\n",
+    "action = MaterialForgeAction(action_type=\"place\", row=3, col=3, atom=\"A\")\n",
+    "obs = env.step(action)\n",
+    "print(f\"\\nAfter placing Metal at (3,3):\")\n",
+    "print(f\"  Reward: {obs.reward}\")\n",
+    "print(f\"  Phase: {obs.phase}\")\n",
+    "print(f\"  Properties: {obs.current_properties}\")\n",
+    "print(\"\\nEnvironment OK!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. TRL Environment Wrapper\n",
+    "\n",
+    "TRL's `GRPOTrainer` with `environment_factory` requires a wrapper class that exposes **named tool methods** (not a generic `step()`). The trainer discovers these methods and presents them to the LLM as callable tools.\n",
+    "\n",
+    "Our wrapper exposes three tools:\n",
+    "- `place_atom(row, col, atom)` \u2014 place a new atom on an empty cell\n",
+    "- `remove_atom(row, col)` \u2014 remove an existing atom\n",
+    "- `replace_atom(row, col, atom)` \u2014 swap an atom for a different species\n",
+    "\n",
+    "Each tool internally calls the OpenEnv environment's standard `step()` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import random\nfrom typing import Optional\n\nfrom environment.config import GRID_SIZE, PROPERTY_NAMES, ATOM_TYPES\nfrom environment.rubrics import HeuristicRewardRubric\nfrom models import MaterialForgeAction, MaterialForgeObservation\nfrom server.material_forge_env_environment import MaterialForgeEnvironment\n\n\nDIFFICULTIES = [\"easy\", \"medium\", \"hard\"]\n\n\nclass MaterialForgeTRLEnv:\n    \"\"\"TRL-compatible wrapper for MaterialForge crystal design.\"\"\"\n\n    def __init__(self):\n        self.env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n        self.reward = 0.0\n        self.done = False\n        self._obs = None\n        self._best_reward = 0.0\n        self._reward_history = []\n        self._invalid_actions = 0\n        self._total_actions = 0\n\n    def reset(self, **kwargs) -> str:\n        seed = random.randint(0, 99999)\n        difficulty = random.choice(DIFFICULTIES)\n        obs = self.env.reset(seed=seed, difficulty=difficulty)\n        self._obs = obs\n        self.reward = 0.0\n        self.done = False\n        self._best_reward = 0.0\n        self._reward_history = []\n        self._invalid_actions = 0\n        self._total_actions = 0\n        return self._format_observation(obs)\n\n    def place_atom(self, row: int, col: int, atom: str) -> str:\n        \"\"\"Place an atom on an empty cell of the crystal lattice.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n            atom: Atom species - 'A' (Metal, hardness, cost 8), 'B' (Conductor, conductivity, cost 6), 'C' (Ceramic, thermal resistance, cost 4), or 'P' (Polymer, elasticity, cost 2)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"place\", row, col, atom)\n\n    def remove_atom(self, row: int, col: int) -> str:\n        \"\"\"Remove an atom from the crystal lattice to free up budget.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"remove\", row, col, None)\n\n    def replace_atom(self, row: int, col: int, atom: str) -> str:\n        \"\"\"Replace an existing atom with a different species to adjust material properties.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n            atom: New atom species - 'A' (Metal), 'B' (Conductor), 'C' (Ceramic), or 'P' (Polymer)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"replace\", row, col, atom)\n\n    def _do_step(self, action_type: str, row: int, col: int, atom: Optional[str]) -> str:\n        if self.done:\n            raise ValueError(\"Episode is over. The crystal design session has ended.\")\n\n        self._total_actions += 1\n        row = max(0, min(int(row), GRID_SIZE - 1))\n        col = max(0, min(int(col), GRID_SIZE - 1))\n\n        if atom is not None:\n            atom = str(atom).upper().strip()\n            if atom not in ATOM_TYPES:\n                atom = \"A\"\n\n        # Pre-check: will this action actually change the grid?\n        grid = self._obs.grid if self._obs else None\n        if grid:\n            cell = grid[row][col]\n            if action_type == \"place\" and cell != \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot place at ({row},{col}) \u2014 cell is already \"\n                    f\"occupied by '{cell}'. Use replace_atom to change it, or pick an empty cell (marked '.').\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"remove\" and cell == \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot remove from ({row},{col}) \u2014 cell is empty. \"\n                    f\"Pick a cell that contains an atom.\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"replace\" and cell == \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot replace at ({row},{col}) \u2014 cell is empty. \"\n                    f\"Use place_atom instead.\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"replace\" and cell == atom:\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cell ({row},{col}) already contains '{atom}'. \"\n                    f\"Pick a different atom species to replace it with.\\n\"\n                    + self._format_observation(self._obs)\n                )\n\n        action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)\n        obs = self.env.step(action)\n        self._obs = obs\n        step_reward = obs.reward if obs.reward is not None else 0.0\n        self._reward_history.append(step_reward)\n        self._best_reward = max(self._best_reward, step_reward)\n        self.reward = self._best_reward\n        self.done = obs.done\n        return self._format_observation(obs)\n\n    def _format_observation(self, obs) -> str:\n        grid_lines = []\n        for r, row in enumerate(obs.grid):\n            grid_lines.append(f\"  {r}: {' '.join(cell if cell != '.' else '.' for cell in row)}\")\n        grid_str = \"\\n\".join(grid_lines)\n\n        gaps = {}\n        for prop in PROPERTY_NAMES:\n            t = obs.target.get(prop, 0.0)\n            c = obs.current_properties.get(prop, 0.0)\n            gaps[prop] = round(t - c, 1)\n\n        props_str = \", \".join(\n            f\"{p}={obs.current_properties.get(p, 0):.1f}/{obs.target.get(p, 0):.1f}(gap:{gaps[p]:+.1f})\"\n            for p in PROPERTY_NAMES\n        )\n\n        sb = obs.score_breakdown or {}\n        stability = sb.get(\"structural_stability\", 0.0)\n        order = sb.get(\"lattice_order_index\", 0.0)\n\n        return (\n            f\"Step {obs.step_number}/{obs.max_steps} | \"\n            f\"Cost: {obs.total_cost:.0f}/{obs.cost_budget:.0f} | \"\n            f\"Phase: {obs.phase} | \"\n            f\"Reward: {obs.reward:.4f}\\n\"\n            f\"Properties: {props_str}\\n\"\n            f\"Stability: {stability:.3f} | Order: {order:.3f}\\n\"\n            f\"Grid (row: cols 0-7):\\n{grid_str}\"\n        )\n\n\nprint(\"MaterialForgeTRLEnv defined.\")"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Test the wrapper"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_env = MaterialForgeTRLEnv()\n",
+    "init_obs = test_env.reset()\n",
+    "print(\"=== Initial State ===\")\n",
+    "print(init_obs)\n",
+    "\n",
+    "print(\"\\n=== After place_atom(3, 3, 'A') ===\")\n",
+    "result = test_env.place_atom(3, 3, \"A\")\n",
+    "print(result)\n",
+    "print(f\"\\nReward stored: {test_env.reward:.4f}\")\n",
+    "print(f\"Done: {test_env.done}\")\n",
+    "\n",
+    "print(\"\\n=== After place_atom(3, 4, 'B') ===\")\n",
+    "result = test_env.place_atom(3, 4, \"B\")\n",
+    "print(result)\n",
+    "print(f\"Reward: {test_env.reward:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Reward Function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def reward_func(environments, **kwargs) -> list[float]:\n    \"\"\"Reward = best_reward + spatial_bonus + phase_bonus - invalid_penalty.\"\"\"\n    rewards = []\n    for env in environments:\n        best = env._best_reward\n        total = max(env._total_actions, 1)\n\n        # Penalty for invalid actions\n        invalid_ratio = env._invalid_actions / total\n        invalid_penalty = 0.3 * invalid_ratio\n\n        # Spatial diversity bonus: reward using multiple rows and columns\n        spatial_bonus = 0.0\n        if env._obs is not None:\n            grid = env._obs.grid\n            rows_used = set()\n            cols_used = set()\n            for r, row in enumerate(grid):\n                for c, cell in enumerate(row):\n                    if cell != \".\":\n                        rows_used.add(r)\n                        cols_used.add(c)\n            # Bonus for spreading across rows AND columns (max 0.10)\n            row_spread = min(len(rows_used) / 4.0, 1.0)  # 4+ rows = full bonus\n            col_spread = min(len(cols_used) / 4.0, 1.0)  # 4+ cols = full bonus\n            spatial_bonus = 0.10 * (row_spread * col_spread)\n\n        # Phase bonus\n        phase_bonus = 0.0\n        if env._obs is not None:\n            phase = env._obs.phase\n            if phase == \"crystalline\":\n                phase_bonus = 0.15\n            elif phase == \"polycrystalline\":\n                phase_bonus = 0.05\n\n        final = best + spatial_bonus + phase_bonus - invalid_penalty\n        rewards.append(max(min(final, 1.0), 0.0))\n    return rewards\n\nprint(\"Reward function: best_reward + spatial_bonus + phase_bonus - invalid_penalty\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Training Dataset\n",
+    "\n",
+    "We create a prompt dataset where each entry instructs the LLM to design a crystal.\n",
+    "The environment generates random targets on each `reset()`, so every episode is unique."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import Dataset\n\nSYSTEM_PROMPT = \"\"\"You are an expert materials scientist designing crystal structures on an 8x8 atomic lattice.\n\nYour goal: place atoms to match target physical properties while staying within budget.\n\nAvailable atom species:\n- A (Metal): Primary contribution to hardness, cost 8 per atom\n- B (Conductor): Primary contribution to conductivity, cost 6 per atom\n- C (Ceramic): Primary contribution to thermal resistance, cost 4 per atom\n- P (Polymer): Primary contribution to elasticity, cost 2 per atom\n\nStrategy:\n- START from the center of the grid (around row 3-4, col 3-4) and expand outward in all directions\n- Build 2D clusters, NOT single rows or columns \u2014 atoms need neighbors in multiple directions for bonding\n- Group same-type atoms together in 2x2 or 3x3 blocks for maximum bonding bonuses\n- After each placement, READ the property gaps and choose the atom type that closes the LARGEST gap\n- Aim for crystalline phase: fill a compact rectangular region, not scattered atoms\n- Stay within the cost budget \u2014 use cheaper atoms (P=2, C=4) when budget is tight\n- Use replace_atom to fix overshooting properties, remove_atom if over budget\n\nUse the tools place_atom, remove_atom, and replace_atom to build your crystal.\"\"\"\n\nUSER_PROMPT = \"Design a crystal lattice that matches the target material properties. Start placing atoms from the center of the grid (row 3-4, col 3-4) and expand outward in all directions. Analyze the property gaps after each action and choose your next atom strategically.\"\n\nNUM_EPISODES = 200\n\ndataset = Dataset.from_dict({\n    \"prompt\": [\n        [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\", \"content\": USER_PROMPT},\n        ]\n    ] * NUM_EPISODES\n})\n\nprint(f\"Dataset created: {len(dataset)} episodes\")\nprint(f\"Sample prompt roles: {[m['role'] for m in dataset[0]['prompt']]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Load Model with Unsloth\n",
+    "\n",
+    "We use Unsloth for memory-efficient 4-bit QLoRA fine-tuning. Starting with `Qwen/Qwen3-0.6B` which fits on a free Colab T4 GPU. For better results with more compute, switch to `Qwen/Qwen3-1.7B` or `Qwen/Qwen2.5-3B-Instruct`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import torch\nfrom unsloth import FastLanguageModel\n\nMODEL_NAME = \"Qwen/Qwen3-0.6B\"  # Change to Qwen3-1.7B or Qwen2.5-3B-Instruct for better results\nMAX_SEQ_LENGTH = 4096\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    MODEL_NAME,\n    max_seq_length=MAX_SEQ_LENGTH,\n    load_in_4bit=True,\n    dtype=None,\n)\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=16,\n    lora_alpha=16,\n    lora_dropout=0,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n    use_gradient_checkpointing=\"unsloth\",\n)\n\n# --- Fix Unsloth + TRL compatibility for Qwen3 tool calling ---\n# Unsloth modifies the chat template, breaking TRL's string-comparison checks.\n# We load TRL's own Qwen3 training template (prefix-preserving, has {% generation %}\n# markers) and response schema directly.\nfrom pathlib import Path\nimport trl.chat_template_utils as _ctu\n\n# 1) Set the response schema so TRL can parse <tool_call> blocks\ntokenizer.response_schema = {\n    \"x-regex\": r\"^(?:<think>\\n?(?:(?P<reasoning_content>.*?\\S.*?)\\n?|[\\s]*)</think>\\s*)?(?P<content>.*?)(?:\\n(?=<tool_call>))?(?=(?:<tool_call>|<\\|im_end\\|>|$))(?P<tool_calls>(?:<tool_call>.+?</tool_call>\\s*)+)?\\s*(?:<\\|im_end\\|>|$)\",\n    \"type\": \"object\",\n    \"properties\": {\n        \"role\": {\"const\": \"assistant\"},\n        \"content\": {\"type\": \"string\"},\n        \"reasoning_content\": {\"type\": \"string\"},\n        \"tool_calls\": {\n            \"type\": \"array\",\n            \"x-regex-iterator\": r\"<tool_call>\\s*(.+?)\\s*</tool_call>\",\n            \"items\": {\n                \"x-parser\": \"json\",\n                \"x-parser-args\": {\"transform\": \"{type: 'function', function: @}\"},\n                \"type\": \"object\",\n                \"properties\": {\n                    \"type\": {\"const\": \"function\"},\n                    \"function\": {\n                        \"type\": \"object\",\n                        \"properties\": {\n                            \"name\": {\"type\": \"string\"},\n                            \"arguments\": {\"type\": \"object\", \"additionalProperties\": {}},\n                        },\n                    },\n                },\n            },\n        },\n    },\n}\n\n# 2) Replace Unsloth's modified template with TRL's training-safe Qwen3 template\n_templates_dir = Path(_ctu.__file__).parent / \"chat_templates\"\ntokenizer.chat_template = (_templates_dir / \"qwen3_training.jinja\").read_text()\n\nprint(f\"Model loaded: {MODEL_NAME}\")\nprint(f\"Trainable params: {model.print_trainable_parameters()}\")\nprint(f\"Chat template: TRL qwen3_training.jinja (prefix-preserving)\")\nprint(f\"Response schema: Qwen3 tool-call parser\")"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. GRPO Training\n",
+    "\n",
+    "The `GRPOTrainer` with `environment_factory` handles the full multi-turn loop:\n",
+    "1. Generates completions from the LLM\n",
+    "2. Parses tool calls (`place_atom`, `remove_atom`, `replace_atom`)\n",
+    "3. Executes them against the MaterialForge environment\n",
+    "4. Feeds observations back to the model\n",
+    "5. Collects rewards and updates the policy via GRPO"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import GRPOConfig, GRPOTrainer\n\nOUTPUT_DIR = \"materialforge-grpo-checkpoints\"\n\ntraining_args = GRPOConfig(\n    output_dir=OUTPUT_DIR,\n\n    # Generation\n    num_generations=4,\n    max_completion_length=2048,\n    temperature=1.0,\n\n    # Optimization\n    learning_rate=5e-5,\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=4,\n    num_train_epochs=1,\n\n    # Logging\n    logging_steps=1,\n    log_completions=True,\n    num_completions_to_print=1,\n\n    # Tool calling\n    chat_template_kwargs={\"enable_thinking\": False},\n\n    # Saving\n    save_steps=50,\n    save_total_limit=2,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_func,\n    train_dataset=dataset,\n    args=training_args,\n    environment_factory=MaterialForgeTRLEnv,\n)\n\nprint(\"Trainer initialized. Starting GRPO training...\")\nprint(f\"  Episodes: {NUM_EPISODES}\")\nprint(f\"  Generations per prompt: {training_args.num_generations}\")\nprint(f\"  Max completion length: {training_args.max_completion_length}\")\nprint(f\"  Temperature: {training_args.temperature}\")\nprint(f\"  Learning rate: {training_args.learning_rate}\")\nprint(f\"  Gradient accumulation: {training_args.gradient_accumulation_steps}\")\nprint(f\"  Total steps: {NUM_EPISODES // (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_result = trainer.train()\n",
+    "print(f\"\\nTraining complete!\")\n",
+    "print(f\"  Total steps: {trainer.state.global_step}\")\n",
+    "print(f\"  Final loss: {train_result.training_loss:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Save Trained Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAVE_DIR = \"materialforge-grpo-trained\"\n",
+    "trainer.save_model(SAVE_DIR)\n",
+    "tokenizer.save_pretrained(SAVE_DIR)\n",
+    "print(f\"Model saved to {SAVE_DIR}/\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 10. Training Reward Curves\n",
+    "\n",
+    "Extract and plot the reward history from training logs. These plots are a **required deliverable** \u2014 they must be saved as images and embedded in the README."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import json\n",
+    "\n",
+    "os.makedirs(\"plots\", exist_ok=True)\n",
+    "\n",
+    "log_history = trainer.state.log_history\n",
+    "\n",
+    "steps = []\n",
+    "rewards = []\n",
+    "losses = []\n",
+    "\n",
+    "for entry in log_history:\n",
+    "    if \"reward\" in entry:\n",
+    "        steps.append(entry.get(\"step\", len(steps)))\n",
+    "        rewards.append(entry[\"reward\"])\n",
+    "    if \"loss\" in entry:\n",
+    "        if len(losses) < len(steps):\n",
+    "            losses.append(entry[\"loss\"])\n",
+    "\n",
+    "# --- Plot 1: Reward over training ---\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "ax.plot(steps[:len(rewards)], rewards, color=\"#4285f4\", linewidth=2, label=\"GRPO Reward\")\n",
+    "if len(rewards) > 5:\n",
+    "    window = max(3, len(rewards) // 10)\n",
+    "    smoothed = [sum(rewards[max(0,i-window):i+1])/len(rewards[max(0,i-window):i+1]) for i in range(len(rewards))]\n",
+    "    ax.plot(steps[:len(smoothed)], smoothed, color=\"#ea4335\", linewidth=2, linestyle=\"--\", label=f\"Smoothed (window={window})\")\n",
+    "ax.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "ax.set_ylabel(\"Reward\", fontsize=12)\n",
+    "ax.set_title(\"MaterialForge GRPO Training \u2014 Reward Curve\", fontsize=14, fontweight=\"bold\")\n",
+    "ax.legend()\n",
+    "ax.grid(True, alpha=0.3)\n",
+    "fig.tight_layout()\n",
+    "fig.savefig(\"plots/reward_curve.png\", dpi=150)\n",
+    "plt.show()\n",
+    "print(\"Saved: plots/reward_curve.png\")\n",
+    "\n",
+    "# --- Plot 2: Loss over training ---\n",
+    "if losses:\n",
+    "    fig2, ax2 = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "    ax2.plot(range(len(losses)), losses, color=\"#34a853\", linewidth=2)\n",
+    "    ax2.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "    ax2.set_ylabel(\"Loss\", fontsize=12)\n",
+    "    ax2.set_title(\"MaterialForge GRPO Training \u2014 Loss Curve\", fontsize=14, fontweight=\"bold\")\n",
+    "    ax2.grid(True, alpha=0.3)\n",
+    "    fig2.tight_layout()\n",
+    "    fig2.savefig(\"plots/loss_curve.png\", dpi=150)\n",
+    "    plt.show()\n",
+    "    print(\"Saved: plots/loss_curve.png\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 11. Baseline Comparison\n",
+    "\n",
+    "Compare the trained agent against a random baseline to demonstrate improvement."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random as rng\n",
+    "\n",
+    "def run_random_baseline(n_episodes=20, max_steps=30):\n",
+    "    \"\"\"Run a random agent as baseline.\"\"\"\n",
+    "    env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n",
+    "    episode_rewards = []\n",
+    "    episode_best = []\n",
+    "\n",
+    "    for ep in range(n_episodes):\n",
+    "        obs = env.reset(seed=ep * 7, difficulty=\"easy\")\n",
+    "        rewards = []\n",
+    "        for step in range(max_steps):\n",
+    "            action_type = rng.choice([\"place\", \"place\", \"place\", \"replace\", \"remove\"])\n",
+    "            row = rng.randint(0, 7)\n",
+    "            col = rng.randint(0, 7)\n",
+    "            atom = rng.choice([\"A\", \"B\", \"C\", \"P\"]) if action_type != \"remove\" else None\n",
+    "            action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)\n",
+    "            obs = env.step(action)\n",
+    "            rewards.append(obs.reward if obs.reward else 0.0)\n",
+    "            if obs.done:\n",
+    "                break\n",
+    "        episode_rewards.append(sum(rewards) / len(rewards) if rewards else 0.0)\n",
+    "        episode_best.append(max(rewards) if rewards else 0.0)\n",
+    "\n",
+    "    return episode_rewards, episode_best\n",
+    "\n",
+    "print(\"Running random baseline (20 episodes)...\")\n",
+    "random_avg_rewards, random_best_rewards = run_random_baseline(20)\n",
+    "print(f\"Random Agent \u2014 Mean reward: {sum(random_avg_rewards)/len(random_avg_rewards):.4f}\")\n",
+    "print(f\"Random Agent \u2014 Mean best:   {sum(random_best_rewards)/len(random_best_rewards):.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Comparison bar chart\n",
+    "trained_mean_reward = sum(rewards) / len(rewards) if rewards else 0.0\n",
+    "random_mean_reward = sum(random_avg_rewards) / len(random_avg_rewards)\n",
+    "\n",
+    "fig3, ax3 = plt.subplots(figsize=(8, 5))\n",
+    "agents = [\"Random\\nBaseline\", \"GRPO-Trained\\nAgent\"]\n",
+    "means = [random_mean_reward, trained_mean_reward]\n",
+    "colors = [\"#ea4335\", \"#4285f4\"]\n",
+    "bars = ax3.bar(agents, means, color=colors, width=0.5, edgecolor=\"white\", linewidth=1.5)\n",
+    "\n",
+    "for bar, val in zip(bars, means):\n",
+    "    ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,\n",
+    "             f\"{val:.4f}\", ha=\"center\", va=\"bottom\", fontweight=\"bold\", fontsize=13)\n",
+    "\n",
+    "ax3.set_ylabel(\"Mean Reward\", fontsize=12)\n",
+    "ax3.set_title(\"MaterialForge \u2014 Random Baseline vs GRPO-Trained Agent\", fontsize=14, fontweight=\"bold\")\n",
+    "ax3.set_ylim(0, max(means) * 1.3 + 0.05)\n",
+    "ax3.grid(axis=\"y\", alpha=0.3)\n",
+    "fig3.tight_layout()\n",
+    "fig3.savefig(\"plots/baseline_comparison.png\", dpi=150)\n",
+    "plt.show()\n",
+    "print(\"Saved: plots/baseline_comparison.png\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 12. Trained Agent Inference Demo\n",
+    "\n",
+    "Run the trained model on a fresh episode to see its crystal design strategy."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "\n",
+    "print(\"Loading trained model for inference...\")\n",
+    "inf_model, inf_tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    SAVE_DIR,\n",
+    "    max_seq_length=MAX_SEQ_LENGTH,\n",
+    "    load_in_4bit=True,\n",
+    ")\n",
+    "FastLanguageModel.for_inference(inf_model)\n",
+    "\n",
+    "demo_env = MaterialForgeTRLEnv()\n",
+    "obs_text = demo_env.reset()\n",
+    "\n",
+    "messages = [\n",
+    "    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "    {\"role\": \"user\", \"content\": USER_PROMPT},\n",
+    "    {\"role\": \"assistant\", \"content\": f\"Let me analyze the initial state:\\n{obs_text}\\n\\nI'll start placing atoms strategically.\"},\n",
+    "]\n",
+    "\n",
+    "print(\"=\" * 60)\n",
+    "print(\"TRAINED AGENT DEMO\")\n",
+    "print(\"=\" * 60)\n",
+    "print(f\"\\nInitial observation:\\n{obs_text}\")\n",
+    "print(f\"\\nAgent designing crystal...\")\n",
+    "print(f\"(Full multi-turn inference requires TRL's tool-call loop.\")\n",
+    "print(f\"See the training logs above for actual agent behavior.)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 13. Summary\n",
+    "\n",
+    "| Metric | Random Baseline | GRPO-Trained Agent |\n",
+    "|--------|----------------|--------------------|\n",
+    "| Mean Reward | *see above* | *see above* |\n",
+    "| Phase Achieved | Amorphous | *see training logs* |\n",
+    "\n",
+    "### Deliverables produced by this notebook:\n",
+    "- `plots/reward_curve.png` \u2014 Reward over training steps\n",
+    "- `plots/loss_curve.png` \u2014 Training loss curve\n",
+    "- `plots/baseline_comparison.png` \u2014 Random vs trained comparison\n",
+    "- `materialforge-grpo-trained/` \u2014 Saved model checkpoint\n",
+    "\n",
+    "### Links\n",
+    "- **Live Environment**: [ArshPathan/material_forge_env](https://huggingface.co/spaces/ArshPathan/material_forge_env)\n",
+    "- **Repository**: [MaterialForge GitHub](https://github.com/Arsh-Pathan/MaterialForge)"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": [],
+   "name": "MaterialForge_GRPO_Training.ipynb"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

training/Traning Runs/Run - III/MaterialForge_GRPO_Training.ipynb

@@ -0,0 +1,490 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# MaterialForge: GRPO Training for Crystal Design\n",
+    "\n",
+    "**Meta PyTorch OpenEnv Hackathon x Scaler School of Technology \u2014 Grand Finale**\n",
+    "\n",
+    "This notebook trains an LLM to autonomously design crystalline structures using reinforcement learning (GRPO) with the MaterialForge environment.\n",
+    "\n",
+    "| Component | Details |\n",
+    "|-----------|----------|\n",
+    "| **Environment** | MaterialForge \u2014 8x8 atomic lattice, 4 atom species, 4 physical properties |\n",
+    "| **Algorithm** | GRPO (Group Relative Policy Optimization) via HuggingFace TRL |\n",
+    "| **Efficiency** | Unsloth for 4-bit QLoRA training |\n",
+    "| **Framework** | OpenEnv-compliant environment |\n",
+    "| **HF Space** | [ArshPathan/material_forge_env](https://huggingface.co/spaces/ArshPathan/material_forge_env) |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Installation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "jupyter": {
+     "is_executing": true
+    }
+   },
+   "source": "%%capture\n!pip install -Uq unsloth\n!pip install -Uq trl>=0.18.0 transformers datasets accelerate peft\n!pip install -Uq \"openenv-core[core]>=0.2.3\" aiofiles openai python-dotenv\n!pip install -Uq matplotlib jmespath",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Clone Environment & Setup Paths\n",
+    "\n",
+    "We clone the MaterialForge repository to access the environment code directly.\n",
+    "The same environment is hosted live at: https://huggingface.co/spaces/ArshPathan/material_forge_env"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "REPO_URL = \"https://huggingface.co/spaces/ArshPathan/material_forge_env\"\n",
+    "REPO_DIR = \"MaterialForge\"\n",
+    "\n",
+    "if not os.path.isdir(REPO_DIR):\n",
+    "    !git clone {REPO_URL} {REPO_DIR}\n",
+    "else:\n",
+    "    print(f\"{REPO_DIR}/ already exists, skipping clone.\")\n",
+    "\n",
+    "import sys\n",
+    "sys.path.insert(0, REPO_DIR)\n",
+    "\n",
+    "print(\"Environment source ready.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Verify Environment Works\n",
+    "\n",
+    "Quick sanity check: reset the environment, take one action, confirm reward is computed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from environment.rubrics import HeuristicRewardRubric\n",
+    "from server.material_forge_env_environment import MaterialForgeEnvironment\n",
+    "from models import MaterialForgeAction\n",
+    "\n",
+    "env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n",
+    "obs = env.reset(seed=42, difficulty=\"easy\")\n",
+    "print(f\"Target properties: {obs.target}\")\n",
+    "print(f\"Cost budget: {obs.cost_budget}\")\n",
+    "print(f\"Max steps: {obs.max_steps}\")\n",
+    "\n",
+    "action = MaterialForgeAction(action_type=\"place\", row=3, col=3, atom=\"A\")\n",
+    "obs = env.step(action)\n",
+    "print(f\"\\nAfter placing Metal at (3,3):\")\n",
+    "print(f\"  Reward: {obs.reward}\")\n",
+    "print(f\"  Phase: {obs.phase}\")\n",
+    "print(f\"  Properties: {obs.current_properties}\")\n",
+    "print(\"\\nEnvironment OK!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. TRL Environment Wrapper\n",
+    "\n",
+    "TRL's `GRPOTrainer` with `environment_factory` requires a wrapper class that exposes **named tool methods** (not a generic `step()`). The trainer discovers these methods and presents them to the LLM as callable tools.\n",
+    "\n",
+    "Our wrapper exposes three tools:\n",
+    "- `place_atom(row, col, atom)` \u2014 place a new atom on an empty cell\n",
+    "- `remove_atom(row, col)` \u2014 remove an existing atom\n",
+    "- `replace_atom(row, col, atom)` \u2014 swap an atom for a different species\n",
+    "\n",
+    "Each tool internally calls the OpenEnv environment's standard `step()` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import random\nfrom typing import Optional\n\nfrom environment.config import GRID_SIZE, PROPERTY_NAMES, ATOM_TYPES\nfrom environment.rubrics import HeuristicRewardRubric\nfrom models import MaterialForgeAction, MaterialForgeObservation\nfrom server.material_forge_env_environment import MaterialForgeEnvironment\n\n\nDIFFICULTIES = [\"easy\", \"medium\", \"hard\"]\n\n\nclass MaterialForgeTRLEnv:\n    \"\"\"TRL-compatible wrapper for MaterialForge crystal design.\"\"\"\n\n    def __init__(self):\n        self.env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n        self.reward = 0.0\n        self.done = False\n        self._obs = None\n        self._best_reward = 0.0\n        self._reward_history = []\n        self._invalid_actions = 0\n        self._total_actions = 0\n\n    def reset(self, **kwargs) -> str:\n        seed = random.randint(0, 99999)\n        difficulty = random.choice(DIFFICULTIES)\n        obs = self.env.reset(seed=seed, difficulty=difficulty)\n        self._obs = obs\n        self.reward = 0.0\n        self.done = False\n        self._best_reward = 0.0\n        self._reward_history = []\n        self._invalid_actions = 0\n        self._total_actions = 0\n        return self._format_observation(obs)\n\n    def place_atom(self, row: int, col: int, atom: str) -> str:\n        \"\"\"Place an atom on an empty cell of the crystal lattice.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n            atom: Atom species - 'A' (Metal, hardness, cost 8), 'B' (Conductor, conductivity, cost 6), 'C' (Ceramic, thermal resistance, cost 4), or 'P' (Polymer, elasticity, cost 2)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"place\", row, col, atom)\n\n    def remove_atom(self, row: int, col: int) -> str:\n        \"\"\"Remove an atom from the crystal lattice to free up budget.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"remove\", row, col, None)\n\n    def replace_atom(self, row: int, col: int, atom: str) -> str:\n        \"\"\"Replace an existing atom with a different species to adjust material properties.\n\n        Args:\n            row: Row index on the 8x8 grid (0-7)\n            col: Column index on the 8x8 grid (0-7)\n            atom: New atom species - 'A' (Metal), 'B' (Conductor), 'C' (Ceramic), or 'P' (Polymer)\n\n        Returns:\n            Updated lattice state with current properties, target gaps, and phase.\n        \"\"\"\n        return self._do_step(\"replace\", row, col, atom)\n\n    def _do_step(self, action_type: str, row: int, col: int, atom: Optional[str]) -> str:\n        if self.done:\n            raise ValueError(\"Episode is over. The crystal design session has ended.\")\n\n        self._total_actions += 1\n        row = max(0, min(int(row), GRID_SIZE - 1))\n        col = max(0, min(int(col), GRID_SIZE - 1))\n\n        if atom is not None:\n            atom = str(atom).upper().strip()\n            if atom not in ATOM_TYPES:\n                atom = \"A\"\n\n        # Pre-check: will this action actually change the grid?\n        grid = self._obs.grid if self._obs else None\n        if grid:\n            cell = grid[row][col]\n            if action_type == \"place\" and cell != \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot place at ({row},{col}) \u2014 cell is already \"\n                    f\"occupied by '{cell}'. Use replace_atom to change it, or pick an empty cell (marked '.').\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"remove\" and cell == \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot remove from ({row},{col}) \u2014 cell is empty. \"\n                    f\"Pick a cell that contains an atom.\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"replace\" and cell == \".\":\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cannot replace at ({row},{col}) \u2014 cell is empty. \"\n                    f\"Use place_atom instead.\\n\"\n                    + self._format_observation(self._obs)\n                )\n            if action_type == \"replace\" and cell == atom:\n                self._invalid_actions += 1\n                return (\n                    f\"INVALID ACTION: Cell ({row},{col}) already contains '{atom}'. \"\n                    f\"Pick a different atom species to replace it with.\\n\"\n                    + self._format_observation(self._obs)\n                )\n\n        action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)\n        obs = self.env.step(action)\n        self._obs = obs\n        step_reward = obs.reward if obs.reward is not None else 0.0\n        self._reward_history.append(step_reward)\n        self._best_reward = max(self._best_reward, step_reward)\n        self.reward = self._best_reward\n        self.done = obs.done\n        return self._format_observation(obs)\n\n    def _format_observation(self, obs) -> str:\n        grid_lines = []\n        for r, row in enumerate(obs.grid):\n            grid_lines.append(f\"  {r}: {' '.join(cell if cell != '.' else '.' for cell in row)}\")\n        grid_str = \"\\n\".join(grid_lines)\n\n        gaps = {}\n        for prop in PROPERTY_NAMES:\n            t = obs.target.get(prop, 0.0)\n            c = obs.current_properties.get(prop, 0.0)\n            gaps[prop] = round(t - c, 1)\n\n        props_str = \", \".join(\n            f\"{p}={obs.current_properties.get(p, 0):.1f}/{obs.target.get(p, 0):.1f}(gap:{gaps[p]:+.1f})\"\n            for p in PROPERTY_NAMES\n        )\n\n        sb = obs.score_breakdown or {}\n        stability = sb.get(\"structural_stability\", 0.0)\n        order = sb.get(\"lattice_order_index\", 0.0)\n\n        return (\n            f\"Step {obs.step_number}/{obs.max_steps} | \"\n            f\"Cost: {obs.total_cost:.0f}/{obs.cost_budget:.0f} | \"\n            f\"Phase: {obs.phase} | \"\n            f\"Reward: {obs.reward:.4f}\\n\"\n            f\"Properties: {props_str}\\n\"\n            f\"Stability: {stability:.3f} | Order: {order:.3f}\\n\"\n            f\"Grid (row: cols 0-7):\\n{grid_str}\"\n        )\n\n\nprint(\"MaterialForgeTRLEnv defined.\")"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Test the wrapper"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_env = MaterialForgeTRLEnv()\n",
+    "init_obs = test_env.reset()\n",
+    "print(\"=== Initial State ===\")\n",
+    "print(init_obs)\n",
+    "\n",
+    "print(\"\\n=== After place_atom(3, 3, 'A') ===\")\n",
+    "result = test_env.place_atom(3, 3, \"A\")\n",
+    "print(result)\n",
+    "print(f\"\\nReward stored: {test_env.reward:.4f}\")\n",
+    "print(f\"Done: {test_env.done}\")\n",
+    "\n",
+    "print(\"\\n=== After place_atom(3, 4, 'B') ===\")\n",
+    "result = test_env.place_atom(3, 4, \"B\")\n",
+    "print(result)\n",
+    "print(f\"Reward: {test_env.reward:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Reward Function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def reward_func(environments, **kwargs) -> list[float]:\n    \"\"\"Reward = best_reward + spatial_bonus + phase_bonus - invalid_penalty.\"\"\"\n    rewards = []\n    for env in environments:\n        best = env._best_reward\n        total = max(env._total_actions, 1)\n\n        # Penalty for invalid actions\n        invalid_ratio = env._invalid_actions / total\n        invalid_penalty = 0.3 * invalid_ratio\n\n        # Spatial diversity bonus: reward using multiple rows and columns\n        spatial_bonus = 0.0\n        if env._obs is not None:\n            grid = env._obs.grid\n            rows_used = set()\n            cols_used = set()\n            for r, row in enumerate(grid):\n                for c, cell in enumerate(row):\n                    if cell != \".\":\n                        rows_used.add(r)\n                        cols_used.add(c)\n            # Bonus for spreading across rows AND columns (max 0.10)\n            row_spread = min(len(rows_used) / 4.0, 1.0)  # 4+ rows = full bonus\n            col_spread = min(len(cols_used) / 4.0, 1.0)  # 4+ cols = full bonus\n            spatial_bonus = 0.10 * (row_spread * col_spread)\n\n        # Phase bonus\n        phase_bonus = 0.0\n        if env._obs is not None:\n            phase = env._obs.phase\n            if phase == \"crystalline\":\n                phase_bonus = 0.15\n            elif phase == \"polycrystalline\":\n                phase_bonus = 0.05\n\n        final = best + spatial_bonus + phase_bonus - invalid_penalty\n        rewards.append(max(min(final, 1.0), 0.0))\n    return rewards\n\nprint(\"Reward function: best_reward + spatial_bonus + phase_bonus - invalid_penalty\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Training Dataset\n",
+    "\n",
+    "We create a prompt dataset where each entry instructs the LLM to design a crystal.\n",
+    "The environment generates random targets on each `reset()`, so every episode is unique."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import Dataset\n\nSYSTEM_PROMPT = \"\"\"You are an expert materials scientist designing crystal structures on an 8x8 atomic lattice.\n\nYour goal: place atoms to match target physical properties while staying within budget.\n\nAvailable atom species:\n- A (Metal): Primary contribution to hardness, cost 8 per atom\n- B (Conductor): Primary contribution to conductivity, cost 6 per atom\n- C (Ceramic): Primary contribution to thermal resistance, cost 4 per atom\n- P (Polymer): Primary contribution to elasticity, cost 2 per atom\n\nStrategy:\n- START from the center of the grid (around row 3-4, col 3-4) and expand outward in all directions\n- Build 2D clusters, NOT single rows or columns \u2014 atoms need neighbors in multiple directions for bonding\n- Group same-type atoms together in 2x2 or 3x3 blocks for maximum bonding bonuses\n- After each placement, READ the property gaps and choose the atom type that closes the LARGEST gap\n- Aim for crystalline phase: fill a compact rectangular region, not scattered atoms\n- Stay within the cost budget \u2014 use cheaper atoms (P=2, C=4) when budget is tight\n- Use replace_atom to fix overshooting properties, remove_atom if over budget\n\nUse the tools place_atom, remove_atom, and replace_atom to build your crystal.\"\"\"\n\nUSER_PROMPT = \"Design a crystal lattice that matches the target material properties. Start placing atoms from the center of the grid (row 3-4, col 3-4) and expand outward in all directions. Analyze the property gaps after each action and choose your next atom strategically.\"\n\nNUM_EPISODES = 200\n\ndataset = Dataset.from_dict({\n    \"prompt\": [\n        [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\", \"content\": USER_PROMPT},\n        ]\n    ] * NUM_EPISODES\n})\n\nprint(f\"Dataset created: {len(dataset)} episodes\")\nprint(f\"Sample prompt roles: {[m['role'] for m in dataset[0]['prompt']]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Load Model with Unsloth\n",
+    "\n",
+    "We use Unsloth for memory-efficient 4-bit QLoRA fine-tuning. Starting with `Qwen/Qwen3-0.6B` which fits on a free Colab T4 GPU. For better results with more compute, switch to `Qwen/Qwen3-1.7B` or `Qwen/Qwen2.5-3B-Instruct`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import torch\nfrom unsloth import FastLanguageModel\n\nMODEL_NAME = \"Qwen/Qwen3-0.6B\"  # Change to Qwen3-1.7B or Qwen2.5-3B-Instruct for better results\nMAX_SEQ_LENGTH = 4096\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    MODEL_NAME,\n    max_seq_length=MAX_SEQ_LENGTH,\n    load_in_4bit=True,\n    dtype=None,\n)\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=16,\n    lora_alpha=16,\n    lora_dropout=0,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                    \"gate_proj\", \"up_proj\", \"down_proj\"],\n    use_gradient_checkpointing=\"unsloth\",\n)\n\n# --- Fix Unsloth + TRL compatibility for Qwen3 tool calling ---\n# Unsloth modifies the chat template, breaking TRL's string-comparison checks.\n# We load TRL's own Qwen3 training template (prefix-preserving, has {% generation %}\n# markers) and response schema directly.\nfrom pathlib import Path\nimport trl.chat_template_utils as _ctu\n\n# 1) Set the response schema so TRL can parse <tool_call> blocks\ntokenizer.response_schema = {\n    \"x-regex\": r\"^(?:<think>\\n?(?:(?P<reasoning_content>.*?\\S.*?)\\n?|[\\s]*)</think>\\s*)?(?P<content>.*?)(?:\\n(?=<tool_call>))?(?=(?:<tool_call>|<\\|im_end\\|>|$))(?P<tool_calls>(?:<tool_call>.+?</tool_call>\\s*)+)?\\s*(?:<\\|im_end\\|>|$)\",\n    \"type\": \"object\",\n    \"properties\": {\n        \"role\": {\"const\": \"assistant\"},\n        \"content\": {\"type\": \"string\"},\n        \"reasoning_content\": {\"type\": \"string\"},\n        \"tool_calls\": {\n            \"type\": \"array\",\n            \"x-regex-iterator\": r\"<tool_call>\\s*(.+?)\\s*</tool_call>\",\n            \"items\": {\n                \"x-parser\": \"json\",\n                \"x-parser-args\": {\"transform\": \"{type: 'function', function: @}\"},\n                \"type\": \"object\",\n                \"properties\": {\n                    \"type\": {\"const\": \"function\"},\n                    \"function\": {\n                        \"type\": \"object\",\n                        \"properties\": {\n                            \"name\": {\"type\": \"string\"},\n                            \"arguments\": {\"type\": \"object\", \"additionalProperties\": {}},\n                        },\n                    },\n                },\n            },\n        },\n    },\n}\n\n# 2) Replace Unsloth's modified template with TRL's training-safe Qwen3 template\n_templates_dir = Path(_ctu.__file__).parent / \"chat_templates\"\ntokenizer.chat_template = (_templates_dir / \"qwen3_training.jinja\").read_text()\n\nprint(f\"Model loaded: {MODEL_NAME}\")\nprint(f\"Trainable params: {model.print_trainable_parameters()}\")\nprint(f\"Chat template: TRL qwen3_training.jinja (prefix-preserving)\")\nprint(f\"Response schema: Qwen3 tool-call parser\")"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. GRPO Training\n",
+    "\n",
+    "The `GRPOTrainer` with `environment_factory` handles the full multi-turn loop:\n",
+    "1. Generates completions from the LLM\n",
+    "2. Parses tool calls (`place_atom`, `remove_atom`, `replace_atom`)\n",
+    "3. Executes them against the MaterialForge environment\n",
+    "4. Feeds observations back to the model\n",
+    "5. Collects rewards and updates the policy via GRPO"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import GRPOConfig, GRPOTrainer\n\nOUTPUT_DIR = \"materialforge-grpo-checkpoints\"\n\ntraining_args = GRPOConfig(\n    output_dir=OUTPUT_DIR,\n\n    # Generation\n    num_generations=4,\n    max_completion_length=2048,\n    temperature=1.0,\n\n    # Optimization\n    learning_rate=5e-5,\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=4,\n    num_train_epochs=1,\n\n    # Logging\n    logging_steps=1,\n    log_completions=True,\n    num_completions_to_print=1,\n\n    # Tool calling\n    chat_template_kwargs={\"enable_thinking\": False},\n\n    # Saving\n    save_steps=50,\n    save_total_limit=2,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_func,\n    train_dataset=dataset,\n    args=training_args,\n    environment_factory=MaterialForgeTRLEnv,\n)\n\nprint(\"Trainer initialized. Starting GRPO training...\")\nprint(f\"  Episodes: {NUM_EPISODES}\")\nprint(f\"  Generations per prompt: {training_args.num_generations}\")\nprint(f\"  Max completion length: {training_args.max_completion_length}\")\nprint(f\"  Temperature: {training_args.temperature}\")\nprint(f\"  Learning rate: {training_args.learning_rate}\")\nprint(f\"  Gradient accumulation: {training_args.gradient_accumulation_steps}\")\nprint(f\"  Total steps: {NUM_EPISODES // (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_result = trainer.train()\n",
+    "print(f\"\\nTraining complete!\")\n",
+    "print(f\"  Total steps: {trainer.state.global_step}\")\n",
+    "print(f\"  Final loss: {train_result.training_loss:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Save Trained Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAVE_DIR = \"materialforge-grpo-trained\"\n",
+    "trainer.save_model(SAVE_DIR)\n",
+    "tokenizer.save_pretrained(SAVE_DIR)\n",
+    "print(f\"Model saved to {SAVE_DIR}/\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 10. Training Reward Curves\n",
+    "\n",
+    "Extract and plot the reward history from training logs. These plots are a **required deliverable** \u2014 they must be saved as images and embedded in the README."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import json\n",
+    "\n",
+    "os.makedirs(\"plots\", exist_ok=True)\n",
+    "\n",
+    "log_history = trainer.state.log_history\n",
+    "\n",
+    "steps = []\n",
+    "rewards = []\n",
+    "losses = []\n",
+    "\n",
+    "for entry in log_history:\n",
+    "    if \"reward\" in entry:\n",
+    "        steps.append(entry.get(\"step\", len(steps)))\n",
+    "        rewards.append(entry[\"reward\"])\n",
+    "    if \"loss\" in entry:\n",
+    "        if len(losses) < len(steps):\n",
+    "            losses.append(entry[\"loss\"])\n",
+    "\n",
+    "# --- Plot 1: Reward over training ---\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "ax.plot(steps[:len(rewards)], rewards, color=\"#4285f4\", linewidth=2, label=\"GRPO Reward\")\n",
+    "if len(rewards) > 5:\n",
+    "    window = max(3, len(rewards) // 10)\n",
+    "    smoothed = [sum(rewards[max(0,i-window):i+1])/len(rewards[max(0,i-window):i+1]) for i in range(len(rewards))]\n",
+    "    ax.plot(steps[:len(smoothed)], smoothed, color=\"#ea4335\", linewidth=2, linestyle=\"--\", label=f\"Smoothed (window={window})\")\n",
+    "ax.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "ax.set_ylabel(\"Reward\", fontsize=12)\n",
+    "ax.set_title(\"MaterialForge GRPO Training \u2014 Reward Curve\", fontsize=14, fontweight=\"bold\")\n",
+    "ax.legend()\n",
+    "ax.grid(True, alpha=0.3)\n",
+    "fig.tight_layout()\n",
+    "fig.savefig(\"plots/reward_curve.png\", dpi=150)\n",
+    "plt.show()\n",
+    "print(\"Saved: plots/reward_curve.png\")\n",
+    "\n",
+    "# --- Plot 2: Loss over training ---\n",
+    "if losses:\n",
+    "    fig2, ax2 = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "    ax2.plot(range(len(losses)), losses, color=\"#34a853\", linewidth=2)\n",
+    "    ax2.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "    ax2.set_ylabel(\"Loss\", fontsize=12)\n",
+    "    ax2.set_title(\"MaterialForge GRPO Training \u2014 Loss Curve\", fontsize=14, fontweight=\"bold\")\n",
+    "    ax2.grid(True, alpha=0.3)\n",
+    "    fig2.tight_layout()\n",
+    "    fig2.savefig(\"plots/loss_curve.png\", dpi=150)\n",
+    "    plt.show()\n",
+    "    print(\"Saved: plots/loss_curve.png\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 11. Baseline Comparison\n",
+    "\n",
+    "Compare the trained agent against a random baseline to demonstrate improvement."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random as rng\n",
+    "\n",
+    "def run_random_baseline(n_episodes=20, max_steps=30):\n",
+    "    \"\"\"Run a random agent as baseline.\"\"\"\n",
+    "    env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())\n",
+    "    episode_rewards = []\n",
+    "    episode_best = []\n",
+    "\n",
+    "    for ep in range(n_episodes):\n",
+    "        obs = env.reset(seed=ep * 7, difficulty=\"easy\")\n",
+    "        rewards = []\n",
+    "        for step in range(max_steps):\n",
+    "            action_type = rng.choice([\"place\", \"place\", \"place\", \"replace\", \"remove\"])\n",
+    "            row = rng.randint(0, 7)\n",
+    "            col = rng.randint(0, 7)\n",
+    "            atom = rng.choice([\"A\", \"B\", \"C\", \"P\"]) if action_type != \"remove\" else None\n",
+    "            action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)\n",
+    "            obs = env.step(action)\n",
+    "            rewards.append(obs.reward if obs.reward else 0.0)\n",
+    "            if obs.done:\n",
+    "                break\n",
+    "        episode_rewards.append(sum(rewards) / len(rewards) if rewards else 0.0)\n",
+    "        episode_best.append(max(rewards) if rewards else 0.0)\n",
+    "\n",
+    "    return episode_rewards, episode_best\n",
+    "\n",
+    "print(\"Running random baseline (20 episodes)...\")\n",
+    "random_avg_rewards, random_best_rewards = run_random_baseline(20)\n",
+    "print(f\"Random Agent \u2014 Mean reward: {sum(random_avg_rewards)/len(random_avg_rewards):.4f}\")\n",
+    "print(f\"Random Agent \u2014 Mean best:   {sum(random_best_rewards)/len(random_best_rewards):.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Comparison bar chart\n",
+    "trained_mean_reward = sum(rewards) / len(rewards) if rewards else 0.0\n",
+    "random_mean_reward = sum(random_avg_rewards) / len(random_avg_rewards)\n",
+    "\n",
+    "fig3, ax3 = plt.subplots(figsize=(8, 5))\n",
+    "agents = [\"Random\\nBaseline\", \"GRPO-Trained\\nAgent\"]\n",
+    "means = [random_mean_reward, trained_mean_reward]\n",
+    "colors = [\"#ea4335\", \"#4285f4\"]\n",
+    "bars = ax3.bar(agents, means, color=colors, width=0.5, edgecolor=\"white\", linewidth=1.5)\n",
+    "\n",
+    "for bar, val in zip(bars, means):\n",
+    "    ax3.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,\n",
+    "             f\"{val:.4f}\", ha=\"center\", va=\"bottom\", fontweight=\"bold\", fontsize=13)\n",
+    "\n",
+    "ax3.set_ylabel(\"Mean Reward\", fontsize=12)\n",
+    "ax3.set_title(\"MaterialForge \u2014 Random Baseline vs GRPO-Trained Agent\", fontsize=14, fontweight=\"bold\")\n",
+    "ax3.set_ylim(0, max(means) * 1.3 + 0.05)\n",
+    "ax3.grid(axis=\"y\", alpha=0.3)\n",
+    "fig3.tight_layout()\n",
+    "fig3.savefig(\"plots/baseline_comparison.png\", dpi=150)\n",
+    "plt.show()\n",
+    "print(\"Saved: plots/baseline_comparison.png\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 12. Trained Agent Inference Demo\n",
+    "\n",
+    "Run the trained model on a fresh episode to see its crystal design strategy."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "\n",
+    "print(\"Loading trained model for inference...\")\n",
+    "inf_model, inf_tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    SAVE_DIR,\n",
+    "    max_seq_length=MAX_SEQ_LENGTH,\n",
+    "    load_in_4bit=True,\n",
+    ")\n",
+    "FastLanguageModel.for_inference(inf_model)\n",
+    "\n",
+    "demo_env = MaterialForgeTRLEnv()\n",
+    "obs_text = demo_env.reset()\n",
+    "\n",
+    "messages = [\n",
+    "    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "    {\"role\": \"user\", \"content\": USER_PROMPT},\n",
+    "    {\"role\": \"assistant\", \"content\": f\"Let me analyze the initial state:\\n{obs_text}\\n\\nI'll start placing atoms strategically.\"},\n",
+    "]\n",
+    "\n",
+    "print(\"=\" * 60)\n",
+    "print(\"TRAINED AGENT DEMO\")\n",
+    "print(\"=\" * 60)\n",
+    "print(f\"\\nInitial observation:\\n{obs_text}\")\n",
+    "print(f\"\\nAgent designing crystal...\")\n",
+    "print(f\"(Full multi-turn inference requires TRL's tool-call loop.\")\n",
+    "print(f\"See the training logs above for actual agent behavior.)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 13. Summary\n",
+    "\n",
+    "| Metric | Random Baseline | GRPO-Trained Agent |\n",
+    "|--------|----------------|--------------------|\n",
+    "| Mean Reward | *see above* | *see above* |\n",
+    "| Phase Achieved | Amorphous | *see training logs* |\n",
+    "\n",
+    "### Deliverables produced by this notebook:\n",
+    "- `plots/reward_curve.png` \u2014 Reward over training steps\n",
+    "- `plots/loss_curve.png` \u2014 Training loss curve\n",
+    "- `plots/baseline_comparison.png` \u2014 Random vs trained comparison\n",
+    "- `materialforge-grpo-trained/` \u2014 Saved model checkpoint\n",
+    "\n",
+    "### Links\n",
+    "- **Live Environment**: [ArshPathan/material_forge_env](https://huggingface.co/spaces/ArshPathan/material_forge_env)\n",
+    "- **Repository**: [MaterialForge GitHub](https://github.com/Arsh-Pathan/MaterialForge)"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": [],
+   "name": "MaterialForge_GRPO_Training.ipynb"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

training/train_env_wrapper.py

@@ -1,198 +1,143 @@
-"""
-MaterialForge TRL Environment Wrapper
-======================================
-Wraps the MaterialForge environment for use with TRL's GRPOTrainer
-via the environment_factory pattern. Exposes crystal lattice operations
-as named tool methods that the LLM discovers and calls during training.
-
-The OpenEnv environment itself keeps its standard step() API unchanged.
-This wrapper translates named tool calls into MaterialForgeAction objects.
-"""
-
-import random
-from typing import Optional
-
-import sys
-import os
-
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
-
-from environment.config import GRID_SIZE, PROPERTY_NAMES, ATOM_TYPES
-from environment.lattice import Lattice
-from environment.physics import classify_phase, compute_structural_stability, compute_lattice_order, estimate_properties
-from environment.rubrics import HeuristicRewardRubric
-from models import MaterialForgeAction, MaterialForgeObservation
-from server.material_forge_env_environment import MaterialForgeEnvironment
-
-
-# --- Curriculum Learning State ---
-# Tracks how many episodes have been started across all trainer instances
-# to slowly ramp up difficulty from Easy -> Medium -> Hard.
-TOTAL_EPISODES_STARTED = 0
-
-def get_curriculum_difficulty() -> str:
-    """Selects difficulty based on global episode progress."""
-    global TOTAL_EPISODES_STARTED
-    if TOTAL_EPISODES_STARTED < 40:
-        difficulty = "easy"
-    elif TOTAL_EPISODES_STARTED < 120:
-        difficulty = "medium"
-    else:
-        difficulty = "hard"
-    
-    TOTAL_EPISODES_STARTED += 1
-    return difficulty
-
-
-class MaterialForgeTRLEnv:
-    """TRL-compatible wrapper for MaterialForge crystal design environment.
-
-    Exposes place_atom, remove_atom, replace_atom as tool methods that
-    GRPOTrainer automatically discovers and presents to the LLM as
-    callable functions during multi-turn rollouts.
-    """
-
-    def __init__(self):
-        self.env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())
-        self.reward = 0.0
-        self.done = False
-        self._obs = None
-        
-        # Performance Tracking
-        self._best_reward = 0.0
-        self._invalid_actions = 0
-        self._total_actions = 0
-
-    def reset(self, **kwargs) -> str:
-        """Resets the environment with a curriculum-based difficulty."""
-        seed = random.randint(0, 99999)
-        
-        # Apply Curriculum Learning
-        difficulty = get_curriculum_difficulty()
-        
-        obs = self.env.reset(seed=seed, difficulty=difficulty)
-        self._obs = obs
-        self.reward = 0.0
-        self.done = False
-        
-        # Reset counters
-        self._best_reward = 0.0
-        self._invalid_actions = 0
-        self._total_actions = 0
-        
-        return self._format_observation(obs)
-
-    def place_atom(self, row: int, col: int, atom: str) -> str:
-        """Place an atom on an empty cell of the crystal lattice.
-
-        Args:
-            row: Row index on the 8x8 grid (0-7)
-            col: Column index on the 8x8 grid (0-7)
-            atom: Atom species to place - 'A' (Metal), 'B' (Conductor), 'C' (Ceramic), or 'P' (Polymer)
-
-        Returns:
-            Updated lattice state or an error message if the placement is invalid.
-        """
-        return self._do_step("place", row, col, atom)
-
-    def remove_atom(self, row: int, col: int) -> str:
-        """Remove an atom from the crystal lattice to free up budget.
-
-        Args:
-            row: Row index on the 8x8 grid (0-7)
-            col: Column index on the 8x8 grid (0-7)
-        """
-        return self._do_step("remove", row, col, None)
-
-    def replace_atom(self, row: int, col: int, atom: str) -> str:
-        """Replace an existing atom with a different species.
-
-        Args:
-            row: Row index on the 8x8 grid (0-7)
-            col: Column index on the 8x8 grid (0-7)
-            atom: New atom species - 'A', 'B', 'C', or 'P'
-        """
-        return self._do_step("replace", row, col, atom)
-
-    def _do_step(self, action_type: str, row: int, col: int, atom: Optional[str]) -> str:
-        if self.done:
-            return "ERROR: Episode is over. Stop calling tools and provide your final answer."
-
-        self._total_actions += 1
-        
-        # Clamp inputs to valid range
-        try:
-            row = max(0, min(int(row), GRID_SIZE - 1))
-            col = max(0, min(int(col), GRID_SIZE - 1))
-        except (ValueError, TypeError):
-            self._invalid_actions += 1
-            return "ERROR: Coordinates must be integers between 0 and 7."
-
-        if atom is not None:
-            atom = str(atom).upper().strip()
-            if atom not in ATOM_TYPES:
-                atom = "A"
-
-        # --- Invalid Action Pre-check ---
-        # Prevents the model from making no-op moves that waste steps and budget.
-        grid = self._obs.grid if self._obs else None
-        if grid:
-            cell = grid[row][col]
-            if action_type == "place" and cell != ".":
-                self._invalid_actions += 1
-                return f"INVALID ACTION: ({row},{col}) is already occupied by '{cell}'. Use replace_atom or pick an empty cell ('.')."
-            
-            if action_type == "remove" and cell == ".":
-                self._invalid_actions += 1
-                return f"INVALID ACTION: Cannot remove from ({row},{col}) because it is already empty."
-            
-            if action_type == "replace" and (cell == "." or cell == atom):
-                self._invalid_actions += 1
-                return f"INVALID ACTION: replace_atom at ({row},{col}) is redundant or invalid. Use place_atom if empty."
-
-        # Execute valid action
-        action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)
-        obs = self.env.step(action)
-        self._obs = obs
-        
-        # Track the best reward seen in this episode for the GRPO signal
-        step_reward = obs.reward if obs.reward is not None else 0.0
-        self._best_reward = max(self._best_reward, step_reward)
-        
-        # Signal to GRPOTrainer (includes a penalty for invalid attempts)
-        penalty = 0.05 * self._invalid_actions
-        self.reward = max(self._best_reward - penalty, 0.0)
-        
-        self.done = obs.done
-        return self._format_observation(obs)
-
-    def _format_observation(self, obs: MaterialForgeObservation) -> str:
-        grid_lines = []
-        for r, row in enumerate(obs.grid):
-            grid_lines.append(f"  {r}: {' '.join(cell if cell != '.' else '.' for cell in row)}")
-        grid_str = "\n".join(grid_lines)
-
-        gaps = {}
-        for prop in PROPERTY_NAMES:
-            t = obs.target.get(prop, 0.0)
-            c = obs.current_properties.get(prop, 0.0)
-            gaps[prop] = round(t - c, 1)
-
-        props_str = ", ".join(
-            f"{p}={obs.current_properties.get(p, 0):.1f}/{obs.target.get(p, 0):.1f} (gap:{gaps[p]:+.1f})"
-            for p in PROPERTY_NAMES
-        )
-
-        sb = obs.score_breakdown or {}
-        stability = sb.get("structural_stability", 0.0)
-        order = sb.get("lattice_order_index", 0.0)
-
-        return (
-            f"Step {obs.step_number}/{obs.max_steps} | "
-            f"Cost: {obs.total_cost:.0f}/{obs.cost_budget:.0f} | "
-            f"Phase: {obs.phase} | "
-            f"Reward: {obs.reward:.4f}\n"
-            f"Properties: {props_str}\n"
-            f"Stability: {stability:.3f} | Order: {order:.3f}\n"
-            f"Grid (row: cols 0-7):\n{grid_str}"
-        )
+"""
+MaterialForge TRL Environment Wrapper
+======================================
+Wraps the MaterialForge environment for use with TRL's GRPOTrainer
+via the environment_factory pattern. Exposes crystal lattice operations
+as named tool methods that the LLM discovers and calls during training.
+"""
+
+import random
+from typing import Optional
+import sys
+import os
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from environment.config import GRID_SIZE, PROPERTY_NAMES, ATOM_TYPES
+from environment.rubrics import HeuristicRewardRubric
+from models import MaterialForgeAction, MaterialForgeObservation
+from server.material_forge_env_environment import MaterialForgeEnvironment
+
+
+# --- Curriculum Learning State ---
+TOTAL_EPISODES_STARTED = 0
+
+def get_curriculum_difficulty() -> str:
+    """Selects difficulty based on global episode progress (scaled for 100 episodes)."""
+    global TOTAL_EPISODES_STARTED
+    if TOTAL_EPISODES_STARTED < 25:
+        difficulty = "easy"
+    elif TOTAL_EPISODES_STARTED < 75:
+        difficulty = "medium"
+    else:
+        difficulty = "hard"
+    
+    TOTAL_EPISODES_STARTED += 1
+    return difficulty
+
+
+class MaterialForgeTRLEnv:
+    """TRL-compatible wrapper for MaterialForge crystal design."""
+
+    def __init__(self):
+        self.env = MaterialForgeEnvironment(rubric=HeuristicRewardRubric())
+        self.reward = 0.0
+        self.done = False
+        self._obs = None
+        
+        # Performance Tracking
+        self._best_reward = 0.0
+        self._invalid_actions = 0
+        self._total_actions = 0
+
+    def reset(self, **kwargs) -> str:
+        """Resets the environment with a curriculum-based difficulty."""
+        seed = random.randint(0, 99999)
+        difficulty = get_curriculum_difficulty()
+        
+        obs = self.env.reset(seed=seed, difficulty=difficulty)
+        self._obs = obs
+        self.reward = 0.0
+        self.done = False
+        
+        # Reset counters
+        self._best_reward = 0.0
+        self._invalid_actions = 0
+        self._total_actions = 0
+        
+        return self._format_observation(obs)
+
+    def place_atom(self, row: int, col: int, atom: str) -> str:
+        """Place an atom on an empty cell of the crystal lattice."""
+        return self._do_step("place", row, col, atom)
+
+    def remove_atom(self, row: int, col: int) -> str:
+        """Remove an atom from the crystal lattice."""
+        return self._do_step("remove", row, col, None)
+
+    def replace_atom(self, row: int, col: int, atom: str) -> str:
+        """Replace an existing atom with a different species."""
+        return self._do_step("replace", row, col, atom)
+
+    def _do_step(self, action_type: str, row: int, col: int, atom: Optional[str]) -> str:
+        if self.done:
+            return "ERROR: Episode is over. Stop calling tools."
+
+        self._total_actions += 1
+        
+        try:
+            row = max(0, min(int(row), GRID_SIZE - 1))
+            col = max(0, min(int(col), GRID_SIZE - 1))
+        except (ValueError, TypeError):
+            self._invalid_actions += 1
+            return "ERROR: Coordinates must be integers 0-7."
+
+        if atom is not None:
+            atom = str(atom).upper().strip()
+            if atom not in ATOM_TYPES:
+                atom = "A"
+
+        # --- Invalid Action Pre-check ---
+        grid = self._obs.grid if self._obs else None
+        if grid:
+            cell = grid[row][col]
+            if action_type == "place" and cell != ".":
+                self._invalid_actions += 1
+                return f"INVALID ACTION: ({row},{col}) is already occupied by '{cell}'. Use replace_atom."
+            
+            if action_type == "remove" and cell == ".":
+                self._invalid_actions += 1
+                return f"INVALID ACTION: Cannot remove from empty cell ({row},{col})."
+            
+            if action_type == "replace" and (cell == "." or cell == atom):
+                self._invalid_actions += 1
+                return f"INVALID ACTION: replace_atom at ({row},{col}) is redundant or invalid."
+
+        action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)
+        obs = self.env.step(action)
+        self._obs = obs
+        
+        # Track best reward and apply penalty for invalid attempts
+        step_reward = obs.reward if obs.reward is not None else 0.0
+        self._best_reward = max(self._best_reward, step_reward)
+        
+        penalty = 0.05 * self._invalid_actions
+        self.reward = max(self._best_reward - penalty, 0.0)
+        
+        self.done = obs.done
+        return self._format_observation(obs)
+
+    def _format_observation(self, obs: MaterialForgeObservation) -> str:
+        grid_lines = []
+        for r, row in enumerate(obs.grid):
+            grid_lines.append(f"  {r}: {' '.join(cell if cell != '.' else '.' for cell in row)}")
+        
+        gaps = ", ".join([f"{p}={obs.current_properties.get(p,0):.1f}/{obs.target.get(p,0):.1f}" for p in PROPERTY_NAMES])
+        sb = obs.score_breakdown or {}
+
+        return (
+            f"Step {obs.step_number}/{obs.max_steps} | Cost: {obs.total_cost:.0f}/{obs.cost_budget:.0f} | Phase: {obs.phase}\n"
+            f"Reward: {obs.reward:.4f} | Properties: {gaps}\n"
+            f"Stability: {sb.get('structural_stability',0):.3f} | Order: {sb.get('lattice_order_index',0):.3f}\n"
+            f"Grid:\n{chr(10).join(grid_lines)}"
+        )