from __future__ import annotations import random from typing import Any, Dict, Optional, Tuple from .agent_interface.formatter import format_observation from .agent_interface.parser import parse_agent_output from .common.serialization import to_jsonable from .episode_tracker import EpisodeTracker from .reward.engine import RewardEngine from .task_manager import TaskManager from .world.state import ActionRecord, WorldState from .world_engine import WorldEngine from .actions.registry import ACTION_REGISTRY from .actions.definitions import validate_required_parameters IMMEDIATE_CATASTROPHE_STEP_PENALTY = -0.4 IMMEDIATE_CATASTROPHE_RAW_PENALTY = 4.0 class PermanenceEnv: def __init__(self, config: Optional[Dict[str, Any]] = None) -> None: self.config = dict(config or {}) # Domain filter: "devtools", "meridian", or None for mixed sampling domain = self.config.get("domain", "devtools") self.task_manager = TaskManager(domain=domain) self.world_engine = WorldEngine() self.reward_engine = RewardEngine() self.episode_tracker = EpisodeTracker() self._current_world_state: Optional[WorldState] = None self._current_task = None self._episode_index = 0 def _select_seed(self, seed: Optional[int]) -> int: if seed is not None: return int(seed) return random.Random(self._episode_index + 17).randint(0, 2**31 - 1) def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None): options = options or {} current_episode_index = self._episode_index selected_seed = self._select_seed(seed) force_task = self.config.get("force_task") or options.get("task_id") difficulty = float(options.get("difficulty", self.config.get("difficulty", 0.5))) task_spec, world_state, sampled_params = self.task_manager.instantiate( current_episode_index, selected_seed, force_task, difficulty=difficulty ) self._current_task = task_spec self._current_world_state = world_state self.episode_tracker.reset(task_spec.task_id, world_state.scenario_id, task_spec.max_steps, task_spec.preservation_targets) self._episode_index += 1 observation = format_observation(world_state=world_state, task=task_spec, step=0) info = to_jsonable( { "episode_index": current_episode_index, "task_id": task_spec.task_id, "scenario_id": world_state.scenario_id, "seed": selected_seed, "difficulty": difficulty, "sampled_params": sampled_params, "max_steps": task_spec.max_steps, "available_actions": task_spec.available_actions, "critical_options": world_state.critical_options, } ) return observation, info def _build_step_info(self, **kwargs) -> Dict[str, Any]: return to_jsonable(kwargs) def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, bool, Dict[str, Any]]: assert self._current_world_state is not None, "Call reset() before step()" assert self._current_task is not None, "Call reset() before step()" current_step = self.episode_tracker.increment_step() parsed = parse_agent_output(action) def finalize(reward: float, error_key: str, parse_error_msgs=None): terminated_by_steps = current_step >= self._current_task.max_steps observation = format_observation( world_state=self._current_world_state, task=self._current_task, step=current_step, parse_error=parse_error_msgs, ) return observation, reward, False, terminated_by_steps, self._build_step_info( error=error_key, step=current_step, task_id=self._current_task.task_id, action_id=parsed.action_id, parse_errors=parse_error_msgs or [], ) if parsed.action_id is None: return finalize(-0.1, "parse_failure", parsed.parse_errors) action_def = ACTION_REGISTRY.get(parsed.action_id) if action_def is None: return finalize( -0.1, "unknown_action", [f"Unknown action '{parsed.action_id}'. Choose from: {', '.join(self._current_task.available_actions)}"], ) if parsed.action_id not in self._current_task.available_actions: return finalize( -0.1, "action_not_in_task", [f"'{parsed.action_id}' not available in {self._current_task.task_id}"], ) validation = validate_required_parameters(action_def, parsed.parameters) if not validation.passed: return finalize(-0.1, "missing_parameter", [validation.failure_message]) if parsed.action_id in self._current_world_state.locked_actions: lock_reason = self._current_world_state.locked_actions.get(parsed.action_id, "prior irreversible action") return finalize(-0.2, "action_locked", [f"'{parsed.action_id}' is locked: {lock_reason}"]) for precondition in action_def.preconditions: try: passed = precondition.fn(self._current_world_state, parsed.parameters) except Exception as error: passed = False precondition = type("Failure", (), {"failure_message": f"Precondition error: {error}"})() if not passed: return finalize(-0.1, "precondition_failed", [precondition.failure_message]) try: actual_r_level = action_def.r_level_fn(self._current_world_state, parsed.parameters) actual_r_level = max(1, min(5, int(actual_r_level))) except Exception: actual_r_level = 2 mutation_params = dict(parsed.parameters) mutation_params["__step"] = current_step mutation_params["__action_id"] = parsed.action_id self.world_engine.apply_consequences(self._current_world_state, action_def.consequences, mutation_params) self.episode_tracker.record_prediction( action_id=parsed.action_id, predicted_r_level=parsed.predicted_r_level, predicted_confidence=parsed.predicted_confidence, actual_r_level=actual_r_level, parameters=parsed.parameters, ) self._current_world_state.append_action_record( ActionRecord( action_id=parsed.action_id, step=current_step, parameters=dict(parsed.parameters), actual_r_level=actual_r_level, predicted_r_level=parsed.predicted_r_level, predicted_confidence=parsed.predicted_confidence, ) ) # Latent / background world dynamics — trust decay, deadline # pressure, board expectation drift. Runs AFTER action consequences # so the agent always sees the post-action state in the next obs. try: from .world.dynamics import apply_latent_dynamics apply_latent_dynamics(self._current_world_state, current_step) except Exception: # Dynamics are additive flavour — never block training on a # malformed state. pass predicted = parsed.predicted_r_level is_catastrophic = actual_r_level == 5 and (predicted is None or predicted <= 2) step_reward = IMMEDIATE_CATASTROPHE_STEP_PENALTY if is_catastrophic else 0.0 pre_applied_catastrophe_raw = IMMEDIATE_CATASTROPHE_RAW_PENALTY if is_catastrophic else 0.0 is_success = self.world_engine.check_success(self._current_world_state, self._current_task) is_max_steps = current_step >= self._current_task.max_steps terminated = is_success or is_catastrophic truncated = is_max_steps and not terminated if terminated or truncated: reason = "success" if is_success else "catastrophic_failure" if is_catastrophic else "max_steps" episode_result = self.episode_tracker.finalize(self._current_world_state, self._current_task, reason) reward_breakdown = self.reward_engine.compute_episode_reward( episode_result, pre_applied_catastrophe_penalty=pre_applied_catastrophe_raw, ) reward = reward_breakdown.total + step_reward info = self._build_step_info( termination_reason=reason, episode_result=episode_result.to_dict(), reward_breakdown=reward_breakdown.to_dict(), step=current_step, task_id=self._current_task.task_id, immediate_step_penalty=step_reward, ) else: reward = step_reward info = self._build_step_info( step=current_step, task_id=self._current_task.task_id, action_id=parsed.action_id, action_r_level=actual_r_level, predicted_r_level=parsed.predicted_r_level, predicted_confidence=parsed.predicted_confidence, immediate_step_penalty=step_reward, ) observation = format_observation(world_state=self._current_world_state, task=self._current_task, step=current_step) return observation, reward, terminated, truncated, info