PLAN.md

MaterialForge — Implementation Plan

AI-Driven Atomic Structure Discovery Environment for the Meta PyTorch OpenEnv Hackathon

---

1. Project Summary

MaterialForge is an OpenEnv-compliant RL environment where an AI agent designs atomic crystal structures to match target material properties (hardness, conductivity, thermal resistance, elasticity). The agent places/replaces/removes atoms on a lattice grid, and the environment scores each configuration against the target using physics-inspired heuristics + optional LLM judging.

Deadline: April 8, 2026 (Round 1 submission)

---

2. Architecture Overview

material_forge_env/
├── __init__.py                          # Public exports
├── models.py                            # Action, Observation, State Pydantic models
├── client.py                            # EnvClient (WebSocket client)
├── config.py                            # Constants: atom types, grid size, property ranges
├── lattice.py                           # Core lattice data structure + manipulation
├── physics.py                           # Property estimation from atomic structure
├── scenarios.py                         # Target profile generators (easy/medium/hard)
├── rubrics.py                           # Reward rubrics (heuristic + optional LLM)
├── llm_utils.py                         # Shared LLM client factory (graceful degradation)
├── hint_engine.py                       # Optional LLM hint system for struggling agents
├── openenv.yaml
├── pyproject.toml
├── README.md
└── server/
    ├── __init__.py
    ├── app.py                           # FastAPI + create_app
    ├── material_forge_env_environment.py # Environment implementation
    ├── Dockerfile
    └── requirements.txt

---

3. Core Data Models

3.1 Atom Types & Properties

Atom	Symbol	Description	Cost	Hardness	Conductivity	Thermal Res	Elasticity
A	A	Strong metal	8	High	Medium	High	Low
B	B	Conductive element	6	Low	High	Low	Medium
C	C	Insulating ceramic	4	Medium	Low	High	Low
P	P	Polymer-like	2	Low	Low	Medium	High
.	empty	Vacant site	0	—	—	—	—

3.2 Action Model (MaterialForgeAction)

class MaterialForgeAction(Action):
    action_type: Literal["place", "replace", "remove"]
    row: int          # 0-indexed grid position
    col: int          # 0-indexed grid position
    atom: Optional[Literal["A", "B", "C", "P"]] = None  # required for place/replace

3.3 Observation Model (MaterialForgeObservation)

class MaterialForgeObservation(Observation):
    grid: List[List[str]]              # Current lattice state (e.g. 8x8)
    target: Dict[str, float]           # Target properties {hardness: 90, ...}
    current_properties: Dict[str, float]  # Estimated properties of current structure
    phase: str                         # "crystalline" | "polycrystalline" | "amorphous"
    total_cost: float                  # Sum of atom costs
    cost_budget: float                 # Max allowed cost
    step_number: int                   # Current step in episode
    max_steps: int                     # Episode length limit
    score_breakdown: Dict[str, float]  # Per-component reward breakdown
    hint: Optional[str] = None         # LLM hint (when available)
    # Inherited: done, reward, metadata

3.4 State

Uses the default OpenEnv State with episode_id and step_count.

---

4. Core Modules

4.1 config.py — Constants & Configuration

• GRID_SIZE = 8 (8×8 lattice)
• MAX_STEPS = 50 per episode
• ATOM_TYPES dict with per-atom property contributions and costs
• PROPERTY_NAMES = ["hardness", "conductivity", "thermal_resistance", "elasticity"]
• COST_BUDGET_DEFAULT = 80
• Difficulty presets: easy (wide tolerance), medium, hard (tight tolerance)

4.2 lattice.py — Lattice Engine

Core data structure representing the material unit cell.

Class Lattice:

• __init__(size: int) — create empty grid
• place(row, col, atom) → bool (validates bounds + empty)
• replace(row, col, atom) → bool (validates bounds + occupied)
• remove(row, col) → bool (validates bounds + occupied)
• get_grid() → List[List[str]]
• count_atoms() → Dict[str, int]
• total_cost() → float
• get_neighbors(row, col) → List[str] — 8-connected neighbors
• clone() → Lattice — deep copy for rollback
• from_grid(grid: List[List[str]]) → Lattice — deserialize

4.3 physics.py — Property Estimation

Heuristic-based physics engine. No real MD simulation, but plausible.

Functions:

• estimate_properties(lattice: Lattice) → Dict[str, float]
  • Hardness: f(metal density, bonding density, packing ratio, crystal order)
  • Conductivity: f(conductive atom fraction, connectivity paths, percolation)
  • Thermal Resistance: f(insulator fraction, structural order, density)
  • Elasticity: f(polymer fraction, amorphous ratio, void fraction)
• classify_phase(lattice: Lattice) → str
  • Checks for repeating patterns → "crystalline"
  • Checks for regional order → "polycrystalline"
  • Default → "amorphous"
• compute_stability(lattice: Lattice) → float
  • Penalizes isolated atoms (no neighbors)
  • Rewards dense bonding clusters
  • Checks structural symmetry

Property calculation strategy: Each property is a weighted sum of atomic contributions + structural bonuses:

property = base_contribution(atom_counts) + neighbor_bonus + phase_bonus + symmetry_bonus

All properties normalized to 0–100 scale.

4.4 scenarios.py — Target Profile Generator

Functions:

• generate_scenario(difficulty: str) → Dict
  • Returns: {target: {hardness, conductivity, ...}, cost_budget, tolerance, max_steps}
  • "easy": large tolerance (±20), generous budget, fewer properties to match
  • "medium": moderate tolerance (±10), standard budget
  • "hard": tight tolerance (±5), tight budget, all properties must match
• generate_llm_scenario() — optional LLM-generated creative scenario with narrative
• Predefined classic scenarios: "diamond-like", "conductor", "heat-shield", "flexible-polymer"

4.5 rubrics.py — Reward System

HeuristicRewardRubric(Rubric):

def forward(action, observation) -> float:
    property_match = 1 - mean(|current[p] - target[p]| / 100 for p in properties)
    stability = compute_stability(lattice)
    phase_bonus = phase_alignment_score(current_phase, target_phase)
    cost_penalty = max(0, (total_cost - budget) / budget)

    reward = (0.50 * property_match
            + 0.25 * stability
            + 0.15 * lattice_quality
            + 0.10 * phase_bonus
            - cost_penalty)
    return clamp(reward, 0.0, 1.0)

MaterialDesignLLMJudge(LLMJudge): (optional, 30% weight when available)

• Prompt: "Given target {target} and structure {grid}, rate quality 0-10"
• Uses HF Inference API via HF_TOKEN
• Falls back to heuristic-only if unavailable

Composite Rubric (via WeightedSum):

• With LLM: WeightedSum([HeuristicRubric, LLMJudge], [0.7, 0.3])
• Without LLM: HeuristicRubric only

4.6 llm_utils.py — LLM Client

• Reads HF_TOKEN from environment
• Creates OpenAIClient pointing to HF Inference API
• Returns None gracefully when no token available
• All LLM features degrade gracefully (heuristic-only mode)

4.7 hint_engine.py — Agent Hints

• Rate-limited: max 3 hints per episode
• Only triggers when agent is struggling (reward declining for 5+ steps)
• Provides textual guidance: "Try adding more metal atoms in the center for hardness"
• Returned in observation.hint field

---

5. Environment Logic (MaterialForgeEnvironment)

5.1 reset()

1. Generate scenario (random difficulty or specified)
2. Initialize empty 8×8 lattice
3. Reset rubric state via self._reset_rubric()
4. Return initial observation with empty grid + target properties

5.2 step(action)

1. Validate action (bounds, atom type, action_type logic)
2. Apply action to lattice (place/replace/remove)
3. Estimate current properties via physics.estimate_properties()
4. Classify structural phase via physics.classify_phase()
5. Compute reward via self._apply_rubric(action, observation)
6. Check termination: done = (step >= max_steps) or (all properties within tolerance)
7. Optionally generate hint if agent is struggling
8. Return observation

5.3 Termination Conditions

• Success: All target properties within tolerance → bonus reward
• Budget exceeded: Cost > budget → penalty, episode continues (agent can remove atoms)
• Max steps: Episode ends, final reward reflects best configuration
• Perfect match: Early termination with maximum reward

---

6. Client (MaterialForgeEnv)

Update the existing scaffold client to:

• Serialize MaterialForgeAction properly
• Parse MaterialForgeObservation with grid, properties, phase, etc.
• Handle the new observation fields in _parse_result

---

7. Deployment

7.1 Local Development

cd material_forge_env
uv sync
uv run server  # http://localhost:8000

7.2 Docker

cd material_forge_env/server
docker build -t material-forge-env:latest .
docker run -p 8000:8000 material-forge-env:latest

7.3 Hugging Face Spaces

cd material_forge_env
openenv push  # Deploys to HF Spaces with Gradio UI

---

8. Implementation Order

Phase 1: Core Engine (Priority)

1. config.py — Define atom types, grid constants, property definitions
2. lattice.py — Grid data structure with place/replace/remove operations
3. physics.py — Property estimation heuristics + phase classification
4. models.py — Rewrite Action/Observation with full MaterialForge fields
5. scenarios.py — Target profile generator with difficulty levels

Phase 2: Environment Integration

6. rubrics.py — HeuristicRewardRubric
7. material_forge_env_environment.py — Wire up reset/step with lattice + physics + rubric
8. client.py — Update client serialization/parsing
9. server/app.py — Update create_app with rubric

Phase 3: LLM Features (Optional Enhancement)

10. llm_utils.py — HF Inference API client
11. rubrics.py — Add LLMJudge + WeightedSum composite
12. hint_engine.py — Contextual hint generation
13. scenarios.py — LLM-generated creative scenarios

Phase 4: Polish & Submission

14. README.md — Full documentation with examples
15. Demo script — demo.py showing agent interaction loop
16. Tests — Validate environment correctness
17. Docker build — Verify containerized deployment
18. HF Spaces — Deploy and get demo link

---

9. Evaluation Criteria Alignment

Criteria	How We Address It
Correctness	Physics heuristics produce plausible material properties; validated with known material archetypes
OpenEnv Compliance	Standard Environment base class, Action/Observation models, create_app, rubric system
Task Design Quality	Rich action space (place/replace/remove × 4 atoms × 64 positions), meaningful search problem
Reward Logic	Multi-component reward: property match + stability + lattice quality + phase alignment - cost
Code Quality	Modular design, type hints, docstrings, separated concerns

---

10. Key Design Decisions

1. 2D grid (not 3D) — Keeps complexity manageable for Round 1; 3D is a Phase 2 extension
2. Heuristic physics — No real DFT/MD simulation; plausible approximations that create a meaningful search landscape
3. 4 atom types — Enough variety for interesting compositions without overwhelming the action space
4. 8×8 grid — 64 positions × 5 options (4 atoms + empty) = rich enough search space
5. 50 max steps — Allows iterative refinement without infinite episodes
6. LLM features optional — Everything works without HF_TOKEN; LLM adds richness when available
7. WeightedSum rubric — Clean separation of heuristic (70%) and LLM (30%) scoring

---

11. Submission Checklist

• Public GitHub repository
• Complete README.md with setup instructions
• requirements.txt in server/
• pyproject.toml with all dependencies
• Demo script (demo.py)
• Working Docker build
• Hugging Face Spaces demo link
• OpenEnv validation passes (openenv validate)