Spaces:
Sleeping
Sleeping
File size: 7,566 Bytes
cd470c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | # MaterialForge Environment Implementation.
# This class acts as the OpenEnv-compliant wrapper for our scientific simulation.
import random
from uuid import uuid4
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State
# Direct imports from root-level modules/packages
from environment.config import (
MIN_ATOMS_FOR_COMPLETION,
PROPERTY_NAMES,
VALID_COMPLETION_PHASES,
)
from environment.lattice import Lattice
from models import MaterialForgeAction, MaterialForgeObservation
from environment.physics import (
classify_phase,
compute_lattice_order,
compute_structural_stability,
estimate_properties,
)
from scenarios.scenarios import generate_scenario
# Main environment class: bridges the abstract OpenEnv protocol with our lattice physics.
class MaterialForgeEnvironment(Environment):
"""RL environment for designing crystal structures on a lattice grid.
The agent places, replaces, or removes atoms on an 8x8 grid to achieve
target material properties (hardness, conductivity, thermal_resistance,
elasticity). Episodes end when max_steps is reached or all properties
are within tolerance of the target.
"""
SUPPORTS_CONCURRENT_SESSIONS: bool = True
def __init__(self, rubric=None):
super().__init__(rubric=rubric)
self._lattice = Lattice()
self._state = State(episode_id=str(uuid4()), step_count=0)
self._scenario = {}
self._target = {}
self._cost_budget = 80.0
self._tolerance = 10.0
self._max_steps = 50
self._difficulty = "medium"
# Reset function: triggers a new scenario and initializes the physics lattice.
def reset(self, seed: int | None = None, **kwargs) -> MaterialForgeObservation:
"""Reset environment with a new scenario.
Args:
seed: Optional random seed for reproducibility.
"""
if seed is not None:
random.seed(seed)
difficulty = kwargs.get("difficulty", "medium")
scenario_name = kwargs.get("scenario_name", None)
self._scenario = generate_scenario(difficulty=difficulty, name=scenario_name)
self._target = self._scenario["target"]
self._cost_budget = float(self._scenario["cost_budget"])
self._tolerance = float(self._scenario["tolerance"])
self._max_steps = int(self._scenario["max_steps"])
self._difficulty = self._scenario["difficulty"]
self._lattice = Lattice()
self._state = State(episode_id=str(uuid4()), step_count=0)
self._reset_rubric()
return self._build_observation(done=False, reward=0.0)
# Core execution step: applies the agent's action and computes the resulting physics state.
def step(self, action: MaterialForgeAction) -> MaterialForgeObservation: # type: ignore[override]
"""Execute one step: apply action, compute properties, score, check done."""
self._state.step_count += 1
# Apply action to lattice
success = self._apply_action(action)
# Compute current state via physics heuristics
properties = estimate_properties(self._lattice)
phase = classify_phase(self._lattice)
stability = compute_structural_stability(self._lattice)
order = compute_lattice_order(self._lattice)
# Check termination
done = self._check_done(properties, phase)
# Build observation with advanced scientific score breakdown
score_breakdown = {
"structural_stability": stability,
"lattice_order_index": order,
"action_valid": 1.0 if success else 0.0,
}
obs = MaterialForgeObservation(
grid=self._lattice.get_grid(),
target=self._target,
current_properties=properties,
phase=phase,
total_cost=self._lattice.total_cost(),
cost_budget=self._cost_budget,
step_number=self._state.step_count,
max_steps=self._max_steps,
score_breakdown=score_breakdown,
done=done,
reward=0.0, # placeholder, rubric will compute
metadata={
"difficulty": self._difficulty,
"scenario_name": self._scenario.get("name"),
"tolerance": self._tolerance,
"action_success": success,
"phase_id": phase,
},
)
# Apply rubric to get reward
obs.reward = self._apply_rubric(action, obs)
return obs
# Internal helper to mutate the lattice state.
def _apply_action(self, action: MaterialForgeAction) -> bool:
"""Apply the action to the lattice. Returns True if action was valid."""
if action.action_type == "place":
if action.atom is None:
return False
return self._lattice.place(action.row, action.col, action.atom)
elif action.action_type == "replace":
if action.atom is None:
return False
return self._lattice.replace(action.row, action.col, action.atom)
elif action.action_type == "remove":
return self._lattice.remove(action.row, action.col)
return False
# Evaluation function: checks if the material has reached a stable crystalline phase.
def _check_done(self, properties: dict, phase: str) -> bool:
"""Check if episode should end based on target matching and structural validity."""
if self._state.step_count >= self._max_steps:
return True
atom_count = self._lattice.atom_count()
# Success requires minimal structure size and ordered phase
if atom_count >= MIN_ATOMS_FOR_COMPLETION and phase in VALID_COMPLETION_PHASES:
all_within = all(
abs(properties.get(p, 0.0) - self._target.get(p, 0.0))
<= self._tolerance
for p in PROPERTY_NAMES
)
if all_within:
return True
return False
# State serialization: packages the current grid and metrics into a standardized object.
def _build_observation(self, done: bool, reward: float) -> MaterialForgeObservation:
"""Build an observation from the current lattice state."""
properties = estimate_properties(self._lattice)
phase = classify_phase(self._lattice)
stability = compute_structural_stability(self._lattice)
order = compute_lattice_order(self._lattice)
return MaterialForgeObservation(
grid=self._lattice.get_grid(),
target=self._target,
current_properties=properties,
phase=phase,
total_cost=self._lattice.total_cost(),
cost_budget=self._cost_budget,
step_number=self._state.step_count,
max_steps=self._max_steps,
score_breakdown={
"structural_stability": stability,
"lattice_order_index": order,
},
done=done,
reward=reward,
metadata={
"difficulty": self._difficulty,
"scenario_name": self._scenario.get("name"),
"tolerance": self._tolerance,
},
)
@property
def state(self) -> State:
return self._state
|