Spaces:
Sleeping
Sleeping
File size: 9,566 Bytes
3c3a3df ff23f69 3c3a3df ff23f69 3c3a3df ff23f69 3c3a3df ff23f69 3c3a3df ff23f69 3c3a3df ff23f69 ae1cf0f ff23f69 ae1cf0f ff23f69 ae1cf0f ff23f69 3c3a3df ff23f69 | 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 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 | """
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
self._best_phase = "amorphous"
self._best_rows_used = 0
self._best_cols_used = 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
self._best_phase = obs.phase
self._best_rows_used = 0
self._best_cols_used = 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 - 'A' (Metal, hardness, cost 8), 'B' (Conductor, conductivity, cost 6), 'C' (Ceramic, thermal resistance, cost 4), or 'P' (Polymer, elasticity, cost 2)
Returns:
Updated lattice state with current properties, target gaps, and phase.
"""
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)
Returns:
Updated lattice state with current properties, target gaps, and phase.
"""
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 to adjust material properties.
Args:
row: Row index on the 8x8 grid (0-7)
col: Column index on the 8x8 grid (0-7)
atom: New atom species - 'A' (Metal), 'B' (Conductor), 'C' (Ceramic), or 'P' (Polymer)
Returns:
Updated lattice state with current properties, target gaps, and phase.
"""
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
self.reward = self._compute_episode_reward()
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
self.reward = self._compute_episode_reward()
return (
f"INVALID ACTION: ({row},{col}) is already occupied by '{cell}'. "
"Use replace_atom.\n" + self._format_observation(self._obs)
)
if action_type == "remove" and cell == ".":
self._invalid_actions += 1
self.reward = self._compute_episode_reward()
return (
f"INVALID ACTION: Cannot remove from empty cell ({row},{col}).\n"
+ self._format_observation(self._obs)
)
if action_type == "replace" and (cell == "." or cell == atom):
self._invalid_actions += 1
self.reward = self._compute_episode_reward()
return (
f"INVALID ACTION: replace_atom at ({row},{col}) is redundant or invalid.\n"
+ self._format_observation(self._obs)
)
action = MaterialForgeAction(action_type=action_type, row=row, col=col, atom=atom)
obs = self.env.step(action)
self._obs = obs
# Cache the structure associated with the best reward so shaping stays coherent.
step_reward = obs.reward if obs.reward is not None else 0.0
if step_reward >= self._best_reward:
self._best_reward = step_reward
self._best_rows_used, self._best_cols_used = self._count_rows_and_cols(obs.grid)
self._best_phase = obs.phase
self.reward = self._compute_episode_reward()
self.done = obs.done
return self._format_observation(obs)
def _count_rows_and_cols(self, grid: list[list[str]]) -> tuple[int, int]:
rows_used = set()
cols_used = set()
for r, row in enumerate(grid):
for c, cell in enumerate(row):
if cell != ".":
rows_used.add(r)
cols_used.add(c)
return len(rows_used), len(cols_used)
def _compute_episode_reward(self) -> float:
# Base property matching
total_reward = self._best_reward
# --- Anti-Collapse & Exploration Bonuses ---
# Reward for actually interacting with the environment (Curiosity)
# We give a bonus for each valid action, up to a limit
action_bonus = min(self._total_actions * 0.02, 0.15)
# Grid Usage Bonus (prevents doing nothing)
# If the grid is empty, give a small negative penalty
structure_bonus = 0.0
if self._best_rows_used > 0:
structure_bonus = 0.10
else:
structure_bonus = -0.10
# --- Spatial & Phase Bonuses (Scientific Quality) ---
row_spread = min(self._best_rows_used / 4.0, 1.0)
col_spread = min(self._best_cols_used / 4.0, 1.0)
spatial_bonus = 0.15 * (row_spread * col_spread)
phase_bonus = 0.0
if self._best_phase == "crystalline":
phase_bonus = 0.20
elif self._best_phase == "polycrystalline":
phase_bonus = 0.10
# --- Penalties ---
# Invalid action penalty
total = max(self._total_actions, 1)
invalid_ratio = self._invalid_actions / total
invalid_penalty = 0.2 * invalid_ratio # Softened to prevent fear-based collapse
final = total_reward + action_bonus + structure_bonus + spatial_bonus + phase_bonus - invalid_penalty
return max(min(final, 1.0), 0.0)
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)}"
)
def reward_func(environments, **kwargs) -> list[float]:
"""Episode reward using the best-scoring state plus aligned shaping penalties."""
return [env._compute_episode_reward() for env in environments]
|