Spaces:
Sleeping
Sleeping
File size: 9,570 Bytes
796da7c | 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 | 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
|