Spaces:
Sleeping
Sleeping
| """ | |
| 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] | |