refactor(core): integrate modular task system for episode management

Introduce a new task-based architecture with abstract ThermalPlantTask class and specific implementations (task1-4) to govern episode disturbances, completions, and policies. Update env/core.py to delegate max_steps, resets, and disturbances to tasks, enhancing modularity and deterministic behavior. Refine transitions.py with improved reward smoothing and initial state clamping for observable gaps. Modify inference.py to embed task descriptions in LLM prompts for better guidance. Add comprehensive unit and integration tests for task functionality.

BREAKING CHANGE: Episode management now requires valid task_id; max_steps are task-specific, altering default episode lengths.

env/core.py

@@ -12,14 +12,16 @@ from env.state import ThermalPlantState
 from env.transitions import (
 	build_coherent_initial_state,
 	integration_step,
+	clamp_state,
 )
 from utils.constants import (
 	ACTION_F_TARGET,
 	ACTION_U_TARGET,
 	DEFAULT_EPISODE_ID,
-	DEFAULT_MAX_STEPS,
 	DEFAULT_TASK_ID,
 )
+from tasks.config import ThermalPlantTask
+from tasks.registry import get_task
 
 
 class ThermalPlantEnv:
@@ -39,13 +41,23 @@ class ThermalPlantEnv:
 
 	def __init__(
 		self,
-		max_steps: int = DEFAULT_MAX_STEPS,
+		max_steps: Optional[int] = None,
 		task_id: str = DEFAULT_TASK_ID,
 		episode_id: int = DEFAULT_EPISODE_ID,
 	) -> None:
-		self.max_steps = max_steps
 		self.task_id = task_id
 		self.episode_id = int(episode_id)
+		# Initialize task
+		try:
+			self._task: ThermalPlantTask = get_task(self.task_id)
+			self.max_steps = max_steps if max_steps is not None else self._task.max_steps
+		except ValueError as exc:
+			logging.warning("__init__(): unknown task_id '%s' - falling back to default '%s' (%s)", self.task_id, DEFAULT_TASK_ID, exc)
+			self.task_id = DEFAULT_TASK_ID
+			self._task = get_task(self.task_id)
+			self.max_steps = max_steps if max_steps is not None else self._task.max_steps
+
+		self._task.reset(self.episode_id)
 		self._state = build_coherent_initial_state(task_id=self.task_id, episode_id=self.episode_id)
 		self._step_count = 0
 
@@ -60,12 +72,18 @@ class ThermalPlantEnv:
 		# fall back to the default task id and log a warning so validator runs
 		# continue (robust behaviour for external callers).
 		try:
+			self._task = get_task(self.task_id)
+			self.max_steps = getattr(self._task, "max_steps", self.max_steps)
 			self._state = build_coherent_initial_state(task_id=self.task_id, episode_id=self.episode_id)
 		except ValueError as exc:
 			logging.warning("reset(): unknown task_id '%s' - falling back to default '%s' (%s)", self.task_id, DEFAULT_TASK_ID, exc)
 			self.task_id = DEFAULT_TASK_ID
+			self._task = get_task(self.task_id)
+			self.max_steps = getattr(self._task, "max_steps", self.max_steps)
 			self._state = build_coherent_initial_state(task_id=self.task_id, episode_id=self.episode_id)
+		
 		self._step_count = 0
+		self._task.reset(self.episode_id)
 		return self._state.to_observation()
 
 	def state(self) -> Dict[str, float]:
@@ -94,6 +112,16 @@ class ThermalPlantEnv:
 			error_msg = "invalid_action_payload"
 			action = {}
 
+		self._step_count += 1
+		
+		# Apply task disturbances before physics
+		deltas, task_event = self._task.apply_disturbance(self._state, self._step_count)
+		if deltas:
+			for k, v in deltas.items():
+				if hasattr(self._state, k):
+					setattr(self._state, k, getattr(self._state, k) + v)
+			self._state = clamp_state(self._state)
+
 		next_state, reward, done_catastrophic, transition_info = integration_step(self._state, action)
 		
 		if invalid_action:
@@ -102,14 +130,15 @@ class ThermalPlantEnv:
 				transition_info["error"] = error_msg
 		
 		self._state = next_state
-		self._step_count += 1
 		
-		done = done_catastrophic or self._step_count >= self.max_steps
+		task_done = hasattr(self._task, "is_completed") and self._task.is_completed(self._state, self._step_count)
+		done = done_catastrophic or task_done or self._step_count >= self.max_steps
 		
 		info: Dict[str, Any] = {
 			"error": transition_info.get("error"),
 			"step_count": self._step_count,
 			"invalid_action": transition_info.get("invalid_action", False),
-			"invalid_action_penalty": transition_info.get("invalid_action_penalty", 0.0)
+			"invalid_action_penalty": transition_info.get("invalid_action_penalty", 0.0),
+			"task_event": task_event
 		}
 		return self._state.to_observation(), float(reward), bool(done), info

env/transitions.py

@@ -218,7 +218,24 @@ def build_coherent_initial_state(task_id: str, episode_id: int) -> ThermalPlantS
 		S_BOUNDS[0],
 		S_BOUNDS[1],
 	)
-	power = clamp(load + regime["G_gap"] - P_STRESS_DRAG * stress, P_BOUNDS[0], P_BOUNDS[1])
+	
+	# Compute gap, allowing it to vary around load for balanced control challenge
+	# The gap can be positive or negative, creating both "reduce" and "increase" scenarios
+	g_gap = regime["G_gap"]
+	
+	# Clamp stress slightly if excessive to prevent artificial P reduction
+	if stress > 0.3:
+		stress = clamp(stress * 0.85, S_BOUNDS[0], S_BOUNDS[1])
+	
+	power = clamp(load + g_gap - P_STRESS_DRAG * stress, P_BOUNDS[0], P_BOUNDS[1])
+	
+	# Ensure observable gap: |P - L| >= 0.04 so LLM sees clear control objective
+	gap_magnitude = abs(power - load)
+	if gap_magnitude < 0.04:
+		# If gap too small, nudge power away from load in direction determined by seed
+		direction = 1.0 if (g_gap >= 0) else -1.0
+		power = clamp(load + direction * 0.08, P_BOUNDS[0], P_BOUNDS[1])
+	
 	control, cooling = _solve_controls(
 		load=load,
 		power=power,
@@ -295,8 +312,8 @@ def integration_step(state: ThermalPlantState, action: Dict[str, float]) -> Tupl
 	def der(x_vec):
 		P, T, Pr, S, D = x_vec
 		
-		# Power: dP/dt = P_GAIN * U * (1 - P/1.0) - P_LOAD_COEF*(P - L) - P_DEG_COEF*P*D
-		dP = C.P_GAIN * next_u * (1.0 - P) - C.P_LOAD_COEF * (P - state.L) - C.P_DEG_COEF * P * D
+		# Power: target tracking for U, minus load drag, minus degradation drag
+		dP = C.P_GAIN * (next_u - P) - C.P_LOAD_COEF * (P - state.L) - C.P_DEG_COEF * P * D
 		
 		# Temperature: dT/dt = a*P^1.1 - b*F^1.05 - c*(T - T_env) - e*D*T
 		# T_env is 0.0 effectively
@@ -372,20 +389,23 @@ def check_catastrophic(state: ThermalPlantState) -> Tuple[bool, str]:
 
 
 def compute_reward(prev_state: ThermalPlantState, state: ThermalPlantState, delta_u: float) -> float:
-	# R_track
-	r_track = -abs(state.P - state.L)
-	
-	# R_safety
-	r_safety = C.SAFETY_PENALTY_COEF * max(0.0, state.T - C.SOFT_T) + C.SAFETY_PENALTY_COEF * max(0.0, state.Pr - C.SOFT_PR)
+	# Tracking reward: smooth linear decay as error increases
+	# At error=0: r_track=1.0, at error=0.2: r_track=0.0, at error=0.4: r_track=-1.0
+	# This avoids the cliff discontinuity and provides smooth gradient for learning
+	tracking_error = abs(state.P - state.L)
+	r_track = 1.0 - tracking_error / 0.2  # Linear decay over 0.2 error window
 	
-	# R_stability
-	r_stability = (delta_u ** 2)
+	# Safety penalty: only penalize when approaching limits
+	safety_violation = max(0.0, state.T - C.SOFT_T) + max(0.0, state.Pr - C.SOFT_PR)
+	r_safety = C.SAFETY_PENALTY_COEF * safety_violation
 	
-	# R_eff
-	r_eff = state.U**2 + state.F**2
+	# Stability penalty: penalize excessive control changes
+	r_stability = C.OSCILLATION_PENALTY_COEF * (delta_u ** 2)
 	
-	# R_invalid is applied outside this by parsing wrapper, but can be added in env core if needed.
+	# Stress penalty: penalize high stress accumulation
+	r_stress = C.STRESS_PENALTY_COEF * state.S if state.S > 0.1 else 0.0
 	
-	reward = C.W_TRACK * r_track - C.W_SAFETY * r_safety - C.W_STABILITY * r_stability - C.W_EFF * r_eff
+	# Combine: tracking reward + safety/stress penalty - control oscillation
+	reward = C.W_TRACK * r_track - C.W_SAFETY * r_safety - C.W_STABILITY * r_stability - 0.2 * r_stress
 	return clamp(reward, -10.0, 10.0)
 

inference.py

@@ -12,6 +12,7 @@ from dotenv import load_dotenv
 load_dotenv()
 
 from env.interface import ConcreteOpenEnvInterface
+from tasks.registry import get_task
 from utils import constants as C
 from utils.logging_utils import canonical_action_string, log_end, log_start, log_step
 from utils.parser import parse_llm_action
@@ -28,18 +29,36 @@ DEFAULT_ACTION = {
 }
 
 # These will be injected by the judge's infrastructure
-HF_TOKEN = os.getenv("HF_TOKEN")
-MODEL_NAME = os.getenv("MODEL_NAME")
-API_BASE_URL = os.getenv("API_BASE_URL")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.3-70B-Instruct")
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 TASK_NAME = os.getenv("THERMAL_PLANT_TASK", C.DEFAULT_TASK_ID)
 EPISODE_ID = int(os.getenv("THERMAL_PLANT_EPISODE_ID", str(C.DEFAULT_EPISODE_ID)))
 DEBUG = os.getenv("DEBUG", "false").lower() in ("1", "true", "yes")
 
+
 SYSTEM_PROMPT = textwrap.dedent(
 	"""
-	You are controlling a thermal plant benchmark. 
-	Return ONLY a compact numeric pair separated by a space representing your chosen `U_target` and `F_target` (e.g. "0.60 0.50").
-	Both targets must be strictly in the range [0, 1]. Do not output JSON, explanations, or prose.
+	You are controlling a thermal plant benchmark.
+	Your objective is to track the required load (L) with your power output (P) while keeping temperature (T) and pressure (Pr) safely below 1.0 to avoid stress (S).
+	
+	State Variables (Inputs):
+	P : Current Power output (track this to L)
+	L : Required Load
+	T : Temperature (keep safely below 1.0)
+	Pr: Pressure (keep safely below 1.0)
+	S : System Stress
+	D : Plant Degradation
+	U : Current Control Valve position (inertia-delayed power target)
+	F : Current Cooling Valve position (inertia-delayed cooling target)
+
+	Action Space (Your Outputs):
+	- U_target: Desired power level [0, 1]. Drives the power valve U.
+	- F_target: Desired cooling level [0, 1]. Drives the cooling valve F.
+
+	Note that your target controls (U_target, F_target) are subject to physical lag. They slowly move the actual valves (U, F) towards your targets. Anticipate delayed responses and system inertia.
+	
+	Return ONLY a compact numeric pair separated by a space representing your chosen `U_target` and `F_target` (e.g. "0.60 0.50"). Do not output JSON, explanations, or prose.
 	"""
 ).strip()
 
@@ -50,6 +69,7 @@ def _format_observation(observation: Dict[str, float]) -> str:
 
 
 def build_user_prompt(
+    task_description: str,
     step: int,
     observation: Dict[str, float],
     last_reward: float,
@@ -59,6 +79,8 @@ def build_user_prompt(
 	history_block = "\n".join(history[-4:]) if history else "None"
 	return textwrap.dedent(
 		f"""
+		Task Objective: {task_description}
+
 		Step: {step}
 		Observation: {_format_observation(observation)}
 		Last reward: {last_reward:.2f}
@@ -70,6 +92,7 @@ def build_user_prompt(
 
 def get_model_response(
     client: OpenAI,
+    task_description: str,
     step: int,
     observation: Dict[str, float],
     last_reward: float,
@@ -81,7 +104,7 @@ def get_model_response(
 			model=MODEL_NAME,
 			messages=[
 				{"role": "system", "content": SYSTEM_PROMPT},
-				{"role": "user", "content": build_user_prompt(step, observation, last_reward, history)},
+				{"role": "user", "content": build_user_prompt(task_description, step, observation, last_reward, history)},
 			],
 			temperature=TEMPERATURE,
 			max_tokens=MAX_TOKENS,
@@ -93,21 +116,21 @@ def get_model_response(
 		return ""
 
 
-def compute_normalized_score(rewards: List[float]) -> float:
+def compute_normalized_score(rewards: List[float], max_steps: int) -> float:
 	"""Convert collected step rewards into a normalized score in [0, 1]."""
 	if not rewards:
 		return 0.0
 	positive_reward = sum(max(0.0, reward) for reward in rewards)
-	return min(max(positive_reward / float(C.DEFAULT_MAX_STEPS), 0.0), 1.0)
+	return min(max(positive_reward / float(max_steps), 0.0), 1.0)
 
 
-def determine_termination_reason(loop_error: Optional[str], done: bool, steps_taken: int) -> str:
+def determine_termination_reason(loop_error: Optional[str], done: bool, steps_taken: int, max_steps: int) -> str:
 	"""Return a stable end-of-episode reason."""
 	if loop_error:
 		return "exception"
 	if done:
 		return "env_done"
-	if steps_taken >= C.DEFAULT_MAX_STEPS:
+	if steps_taken >= max_steps:
 		return "max_steps"
 	return "stopped"
 
@@ -115,10 +138,11 @@ def determine_termination_reason(loop_error: Optional[str], done: bool, steps_ta
 def main() -> None:
 	"""Run a full deterministic inference episode and always emit end logs."""
 	model_name_for_logs = MODEL_NAME or "unset"
-	env = ConcreteOpenEnvInterface(max_steps=C.DEFAULT_MAX_STEPS)
+	env = ConcreteOpenEnvInterface()
 	client: Optional[OpenAI] = None
 	last_valid_action: Optional[Dict[str, float]] = None
 	observation = env.reset(task_id=TASK_NAME, episode_id=EPISODE_ID)
+	max_steps = getattr(env._env, "max_steps", 12) if hasattr(env, "_env") else 12
 	history: List[str] = []
 	rewards: List[float] = []
 	score = 0.0
@@ -131,6 +155,12 @@ def main() -> None:
 
 	log_start(task=TASK_NAME, env=BENCHMARK, model=model_name_for_logs)
 
+	try:
+		task_obj = get_task(TASK_NAME)
+		task_description = getattr(task_obj, "description", "Optimal Thermal Plant Control")
+	except Exception:
+		task_description = "Optimal Thermal Plant Control"
+
 	try:
 		if HF_TOKEN and MODEL_NAME and API_BASE_URL:
 			client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
@@ -141,11 +171,12 @@ def main() -> None:
 				flush=True,
 			)
 
-		for step in range(1, C.DEFAULT_MAX_STEPS + 1):
+		for step in range(1, max_steps + 1):
 			raw_response = ""
 			if client is not None:
 				raw_response = get_model_response(
 					client=client,
+					task_description=task_description,
 					step=step,
 					observation=observation,
 					last_reward=last_reward,
@@ -250,9 +281,9 @@ def main() -> None:
 	except Exception as exc:
 		loop_error = str(exc)
 	finally:
-		score = compute_normalized_score(rewards)
+		score = compute_normalized_score(rewards, max_steps)
 		success = score >= SUCCESS_SCORE_THRESHOLD
-		termination_reason = determine_termination_reason(loop_error=loop_error, done=done, steps_taken=steps_taken)
+		termination_reason = determine_termination_reason(loop_error=loop_error, done=done, steps_taken=steps_taken, max_steps=max_steps)
 		trajectory.summary = TrajectorySummary(
 			task=TASK_NAME,
 			benchmark=BENCHMARK,
@@ -267,4 +298,6 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+	if HF_TOKEN is None:
+		raise ValueError("HF_TOKEN environment variable is required")
+	main()

tasks/config.py

@@ -0,0 +1,58 @@
+"""Task configuration, interfaces, and shared types."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Tuple, Any
+
+from env.state import ThermalPlantState
+
+
+class AgentPolicy(ABC):
+    """A deterministic baseline controller for sanity checks and grading."""
+    
+    @abstractmethod
+    def get_action(self, observation: Dict[str, float]) -> Dict[str, float]:
+        """Return a deterministic control action based on the state observation."""
+        pass
+
+
+class ThermalPlantTask(ABC):
+    """Base class for tasks governing the episodes in the standard benchmark.
+    
+    Tasks control deterministic disturbances over time and provide metadata
+    for graders and OpenEnv integration.
+    """
+
+    task_id: str
+    name: str
+    description: str
+    max_steps: int
+
+    @abstractmethod
+    def reset(self, episode_id: int) -> None:
+        """Reset the task-internal state using the episode ID as a deterministic seed."""
+        pass
+
+    @abstractmethod
+    def apply_disturbance(self, state: ThermalPlantState, step: int) -> Tuple[Dict[str, float], Optional[Dict[str, Any]]]:
+        """Compute relative disturbances for the current step.
+        
+        Args:
+            state: The fully specified current actual state (un-rounded).
+            step: The current 1-indexed environment step.
+            
+        Returns:
+            deltas: A dictionary mapping state keys (e.g., 'L', 'D', 'T') to absolute Delta values
+                    to add to the current state fields. Amplitudes should be derived relative
+                    to the variable domain bounds implicitly.
+            event: An optional dictionary describing the event that just occurred for info logging.
+        """
+        pass
+
+    def is_completed(self, state: ThermalPlantState, step_count: int) -> bool:
+        """Check if the task has successfully reached a terminal stabilizing state."""
+        return False
+
+    @abstractmethod
+    def get_baseline_policy(self) -> AgentPolicy:
+        """Return a simple heuristic or rule-based agent capable of surviving the task."""
+        pass

tasks/registry.py

@@ -0,0 +1,37 @@
+"""Central task discovery and instantiation."""
+
+from typing import Dict, Type, Optional
+
+from tasks.config import ThermalPlantTask
+from tasks.task1 import Task1
+from tasks.task2 import Task2
+from tasks.task3 import Task3
+from tasks.task4 import Task4
+
+
+_TASK_CLASSES: Dict[str, Type[ThermalPlantTask]] = {
+    "task1": Task1,
+    "task2": Task2,
+    "task3": Task3,
+    "task4": Task4,
+}
+
+
+def get_task(task_id: str) -> ThermalPlantTask:
+    """Instantiate and return the canonical task."""
+    cls = _TASK_CLASSES.get(task_id)
+    if not cls:
+        raise ValueError(f"Unknown task_id '{task_id}'. Expected one of {list(_TASK_CLASSES.keys())}")
+    return cls()
+
+
+def task_registry() -> Dict[str, dict]:
+    """Expose registry metadata required by OpenEnv validation."""
+    return {
+        task_id: {
+            "name": getattr(cls, "name", task_id),
+            "description": getattr(cls, "description", ""),
+            "max_steps": getattr(cls, "max_steps", None),
+        }
+        for task_id, cls in _TASK_CLASSES.items()
+    }

tasks/task1.py

@@ -0,0 +1,56 @@
+"""Task 1: Stable Baseline Operation.
+
+Maintain P close to L while staying safe and smooth. Constant load pattern L_t = 0.6.
+"""
+
+from typing import Dict, Optional, Tuple, Any
+
+from env.state import ThermalPlantState
+from tasks.config import ThermalPlantTask, AgentPolicy
+from utils.constants import U_BOUNDS, F_BOUNDS, P_BOUNDS, L_BOUNDS, T_BOUNDS, PR_BOUNDS
+
+class BaselinePolicy(AgentPolicy):
+    """A simple rule-based tracker."""
+    
+    def get_action(self, observation: Dict[str, float]) -> Dict[str, float]:
+        # Track power to load
+        u_target = 0.6
+        if observation["P"] < observation["L"] - 0.05:
+            u_target = min(0.9, observation["U"] + 0.1)
+        elif observation["P"] > observation["L"] + 0.05:
+            u_target = max(0.1, observation["U"] - 0.1)
+            
+        # Heavy cooling if temp/pressure gets hot
+        f_target = 0.4
+        if observation["T"] > 0.85 or observation["Pr"] > 0.85:
+            f_target = 0.8
+            u_target = max(0.1, u_target - 0.2)
+            
+        return {"U_target": u_target, "F_target": f_target}
+
+
+class Task1(ThermalPlantTask):
+    task_id = "task1"
+    name = "Stable Baseline Operation"
+    description = "Maintain power close to a constant load of 0.6 while staying safe and smooth."
+    max_steps = 12
+
+    def reset(self, episode_id: int) -> None:
+        self._seed = episode_id
+
+    def apply_disturbance(self, state: ThermalPlantState, step: int) -> Tuple[Dict[str, float], Optional[Dict[str, Any]]]:
+        deltas = {}
+        # Force L_t = 0.6
+        target_L = 0.6
+        if abs(state.L - target_L) > 1e-5:
+            deltas["L"] = target_L - state.L
+        return deltas, {"type": "constant_load", "target_L": target_L}
+
+    def is_completed(self, state: ThermalPlantState, step_count: int) -> bool:
+        # User requested early stop when P=L.
+        tracking_error = abs(state.P - 0.6)
+        # End task immediately if error is very small (converged to load target)
+        return tracking_error <= 0.02 and step_count >= 3
+
+    def get_baseline_policy(self) -> AgentPolicy:
+        return BaselinePolicy()

tasks/task2.py

@@ -0,0 +1,63 @@
+"""Task 2: Load Following.
+
+Track changing load with minimal lag and overshoot.
+Step changes: t=1..3: L=0.5, t=4..6: L=0.8, t=7..N: L=0.6.
+"""
+
+from typing import Dict, Optional, Tuple, Any
+
+from env.state import ThermalPlantState
+from tasks.config import ThermalPlantTask, AgentPolicy
+from utils.constants import L_BOUNDS, D_BOUNDS, TASK_CODE
+
+class PeriodicPolicy(AgentPolicy):
+    """Tracker robust to periodic disturbances."""
+    
+    def get_action(self, observation: Dict[str, float]) -> Dict[str, float]:
+        # Track power to load aggressively to minimize lag
+        u_target = min(max(observation["U"] + 0.4 * (observation["L"] - observation["P"]), 0.1), 0.9)
+        f_target = 0.5
+        if observation["T"] > 0.8:
+            f_target = 0.75
+            u_target -= 0.15
+        elif observation["T"] > 0.9:
+            f_target = 0.95
+            u_target = 0.1
+            
+        return {"U_target": max(0.0, u_target), "F_target": f_target}
+
+
+class Task2(ThermalPlantTask):
+    task_id = "task2"
+    name = "Load Following"
+    description = "Track step changes in required load (0.5 -> 0.8 -> 0.6) with minimal lag and overshoot."
+    max_steps = 12
+
+    def reset(self, episode_id: int) -> None:
+        self._seed = episode_id
+
+    def apply_disturbance(self, state: ThermalPlantState, step: int) -> Tuple[Dict[str, float], Optional[Dict[str, Any]]]:
+        deltas = {}
+        
+        if step <= 3:
+            target_L = 0.5
+        elif step <= 6:
+            target_L = 0.8
+        else:
+            target_L = 0.6
+            
+        if abs(state.L - target_L) > 1e-5:
+            deltas["L"] = target_L - state.L
+            
+        event = {"type": "load_step", "target_L": target_L}
+        return deltas, event
+
+    def is_completed(self, state: ThermalPlantState, step_count: int) -> bool:
+        # Task 2 has a final load step change at step 7.
+        # So we only allow early completion after the final change has stabilized.
+        if step_count < 8:
+            return False
+        return abs(state.P - 0.6) <= 0.02
+
+    def get_baseline_policy(self) -> AgentPolicy:
+        return PeriodicPolicy()

tasks/task3.py

@@ -0,0 +1,64 @@
+"""Task 3: Preemptive Constraint Management.
+
+L_t = moderately high constant (0.7).
+Stress accumulates when T > 0.9. Agent must reduce risk BEFORE visible failure.
+"""
+
+from typing import Dict, Optional, Tuple, Any
+
+from env.state import ThermalPlantState
+from tasks.config import ThermalPlantTask, AgentPolicy
+from utils.constants import L_BOUNDS, TASK_CODE
+
+class RampPolicy(AgentPolicy):
+    """Proportional tracker."""
+    
+    def get_action(self, observation: Dict[str, float]) -> Dict[str, float]:
+        # Track power to load proportionally
+        u_target = min(observation["U"] + 0.4 * (observation["L"] - observation["P"]), 1.0)
+        u_target = max(u_target, 0.0)
+        
+        f_target = max(0.4, 0.8 * observation["T"])
+        # Preemptively keep T < 0.9 by sharply increasing cooling when safely near the threshold
+        if observation["T"] > 0.85:
+            f_target = 0.95
+            u_target = max(0.0, u_target - 0.2)
+            
+        return {"U_target": u_target, "F_target": f_target}
+
+
+class Task3(ThermalPlantTask):
+    task_id = "task3"
+    name = "Preemptive Constraint Management"
+    description = "Manage moderately high load (0.7) and prevent stress accumulation by keeping Temperature below 0.9 under restricted coolant conditions."
+    max_steps = 12
+
+    def reset(self, episode_id: int) -> None:
+        self._seed = int(episode_id)
+
+    def apply_disturbance(self, state: ThermalPlantState, step: int) -> Tuple[Dict[str, float], Optional[Dict[str, Any]]]:
+        deltas: Dict[str, float] = {}
+        event: Dict[str, Any] = {"type": "constraint_management"}
+
+        target_L = 0.7
+        if abs(state.L - target_L) > 1e-5:
+            deltas["L"] = target_L - state.L
+
+        # Exogenous heat: acts as a coolant deficiency.
+        # Max cooling (F=1.0) is no longer strong enough to overcome both this and P=0.7.
+        deltas["T"] = 0.045
+
+        # Stress accumulates when T > 0.9
+        if state.T > 0.9:
+            deltas["S"] = 0.05 # Additive stress accumulation per step
+            event["stress_warning"] = True
+
+        return deltas, event
+
+    def is_completed(self, state: ThermalPlantState, step_count: int) -> bool:
+        # End task if stable and constraints are well managed.
+        error = abs(state.P - 0.7)
+        return error <= 0.02 and state.T < 0.85 and step_count >= 3
+
+    def get_baseline_policy(self) -> AgentPolicy:
+        return RampPolicy()

tasks/task4.py

@@ -0,0 +1,63 @@
+"""Task 4: Fault Recovery with Degradation.
+
+L=0.6 mostly. Disturbance at t=4: T spike (+0.3).
+Degradation reduces cooling effectiveness over time.
+"""
+
+from typing import Dict, Optional, Tuple, Any
+
+from env.state import ThermalPlantState
+from tasks.config import ThermalPlantTask, AgentPolicy
+from utils.constants import T_BOUNDS, PR_BOUNDS, TASK_CODE
+
+class ShockPolicy(AgentPolicy):
+    """Heavy cooling strategy."""
+    
+    def get_action(self, observation: Dict[str, float]) -> Dict[str, float]:
+        # Track power to load cautiously
+        u_target = min(observation["U"] + 0.2 * (observation["L"] - observation["P"]), 0.8)
+        
+        # Heavy cooling
+        f_target = 0.6
+        if observation["T"] > 0.8 or observation["Pr"] > 0.8:
+            f_target = 1.0
+            u_target = 0.0 # Emergency shutdown
+            
+        return {"U_target": max(0.0, u_target), "F_target": f_target}
+
+
+class Task4(ThermalPlantTask):
+    task_id = "task4"
+    name = "Fault Recovery with Degradation"
+    description = "Recover from a major thermal spike at step 4 while managing ongoing system degradation."
+    max_steps = 12
+
+    def reset(self, episode_id: int) -> None:
+        self._seed = int(episode_id)
+
+    def apply_disturbance(self, state: ThermalPlantState, step: int) -> Tuple[Dict[str, float], Optional[Dict[str, Any]]]:
+        deltas = {}
+        event = {"type": "constant_load"}
+
+        target_L = 0.6
+        if abs(state.L - target_L) > 1e-5:
+            deltas["L"] = target_L - state.L
+
+        if step == 4:
+            deltas["T"] = 0.5
+            event = {"type": "thermal_fault", "T_delta": deltas["T"]}
+            
+        # Ongoing degradation to simulate reduced cooling effectiveness
+        deltas["D"] = 0.02
+
+        return deltas, event
+
+    def is_completed(self, state: ThermalPlantState, step_count: int) -> bool:
+        # Task 4 has a major fault at step 4.
+        # Don't end early before the fault has been processed and stabilized.
+        if step_count < 6:
+            return False
+        return abs(state.P - 0.6) <= 0.02 and state.T < 0.8 and state.Pr < 0.8
+
+    def get_baseline_policy(self) -> AgentPolicy:
+        return ShockPolicy()

tests/adversarial/test_malformed_llm_outputs.py

@@ -20,7 +20,8 @@ def mock_env_vars():
 
 def test_inference_main_with_malformed_outputs(capsys, monkeypatch, mock_env_vars):
     # Mock settings to create a short run
-    monkeypatch.setattr(inference.C, "DEFAULT_MAX_STEPS", 4)
+    import tasks.task2
+    monkeypatch.setattr(tasks.task2.Task2, "max_steps", 4)
     monkeypatch.setattr(inference.C, "INCLUDE_PARSE_ERROR_IN_STEP", True)
     
     # We will simulate exactly 4 steps of LLM responses:

tests/integration/test_inference_loop.py

@@ -20,6 +20,7 @@ MOCK_RESPONSES = [
 
 def mock_get_model_response(
     client: MagicMock,
+    task_description: str,
     step: int,
     observation: Dict[str, float],
     last_reward: float,
@@ -35,7 +36,9 @@ def mock_inference_env(monkeypatch):
     monkeypatch.setenv("MODEL_NAME", "fake_model")
     monkeypatch.setenv("API_BASE_URL", "http://fake.api")
     monkeypatch.setenv("THERMAL_PLANT_EPISODE_ID", "1")
-    monkeypatch.setattr(inference.C, "DEFAULT_MAX_STEPS", 5) # Short episode
+    # Tell the default task it's a short episode
+    import tasks.task2
+    monkeypatch.setattr(tasks.task2.Task2, "max_steps", 5)
     yield
 
 def test_inference_main_prints_exact_stdout_format(mock_inference_env, monkeypatch):

tests/integration/test_tasks_integration.py

@@ -0,0 +1,46 @@
+import pytest
+import io
+import contextlib
+import re
+
+import inference
+from env.core import ThermalPlantEnv
+from utils.constants import TASK_CODE
+
+@pytest.fixture
+def mock_inference_env_tasks(monkeypatch):
+    monkeypatch.setenv("HF_TOKEN", "fake_token")
+    monkeypatch.setenv("MODEL_NAME", "fake_model")
+    monkeypatch.setenv("API_BASE_URL", "http://fake.api")
+    monkeypatch.setenv("THERMAL_PLANT_EPISODE_ID", "42")
+    yield monkeypatch
+
+
+@pytest.mark.parametrize("task_id", ["task1", "task2", "task3", "task4"])
+def test_inference_runs_full_episode_for_each_task(mock_inference_env_tasks, task_id):
+    mock_inference_env_tasks.setenv("THERMAL_PLANT_TASK", task_id)
+    
+    # We will hook inference so it just runs the task's baseline policy instead of calling a model
+    env = ThermalPlantEnv(task_id=task_id, episode_id=42)
+    policy = env._task.get_baseline_policy()
+
+    def mock_get_model_response(client, task_description, step, observation, last_reward, history):
+        action = policy.get_action(observation)
+        # convert back to valid string
+        return f"{action['U_target']:.2f} {action['F_target']:.2f}"
+    
+    mock_inference_env_tasks.setattr(inference, "get_model_response", mock_get_model_response)
+
+    stdout_capture = io.StringIO()
+    with contextlib.redirect_stdout(stdout_capture):
+        inference.main()
+        
+    output = stdout_capture.getvalue()
+    lines = output.strip().split("\n")
+    
+    assert lines[0].startswith(f"[START] task={task_id} env=thermal-plant-control")
+    assert lines[-1].startswith("[END] success=")
+    
+    step_lines = [line for line in lines if line.startswith("[STEP]")]
+    max_steps = env.max_steps
+    assert len(step_lines) <= max_steps

tests/unit/test_tasks.py

@@ -0,0 +1,89 @@
+import pytest
+
+from env.core import ThermalPlantEnv
+from tasks.registry import get_task
+
+
+@pytest.mark.parametrize("task_id", ["task1", "task2", "task3", "task4"])
+def test_task_registry_has_tasks(task_id):
+    task = get_task(task_id)
+    assert task.task_id == task_id
+    assert task.max_steps > 0
+    assert task.get_baseline_policy() is not None
+
+
+def test_task1_baseline():
+    env = ThermalPlantEnv(task_id="task1", episode_id=42)
+    obs = env.reset()
+    assert env.max_steps == 12
+
+    # Task1 applies NO disturbances except constant Load enforcement
+    action = env._task.get_baseline_policy().get_action(obs)
+    obs, reward, done, info = env.step(action)
+    
+    assert info.get("task_event") is not None
+    assert info["task_event"]["type"] == "constant_load"
+
+
+def test_task2_periodic_pulses():
+    env = ThermalPlantEnv(task_id="task2", episode_id=42)
+    obs = env.reset()
+    assert env.max_steps == 12
+    
+    events_seen = 0
+    for step in range(1, 13):
+        action = env._task.get_baseline_policy().get_action(obs)
+        obs, reward, done, info = env.step(action)
+        if info.get("task_event"):
+            assert info["task_event"]["type"] == "load_step"
+            target = info["task_event"]["target_L"]
+            if step <= 3:
+                assert target == 0.5
+            elif step <= 6:
+                assert target == 0.8
+            else:
+                assert target == 0.6
+                
+        if done:
+            break
+
+
+def test_task3_ramp():
+    env = ThermalPlantEnv(task_id="task3", episode_id=42)
+    obs = env.reset()
+    assert env.max_steps == 12
+    
+    events_seen = 0
+    # Artificially spike T to test stress accumulation
+    env._state.T = 0.95
+    for _ in range(12):
+        action = env._task.get_baseline_policy().get_action(obs)
+        obs, reward, done, info = env.step(action)
+        if info.get("task_event") is not None:
+            events_seen += 1
+            
+        if done:
+            break
+            
+    assert events_seen > 0, "No events occurred in task3"
+
+
+def test_task4_thermal_shock():
+    env = ThermalPlantEnv(task_id="task4", episode_id=42)
+    obs = env.reset()
+    assert env.max_steps == 12
+    
+    events_seen = 0
+    for step in range(1, 13):
+        action = env._task.get_baseline_policy().get_action(obs)
+        obs, reward, done, info = env.step(action)
+        if info.get("task_event") is not None:
+            if info["task_event"].get("type") == "thermal_fault":
+                events_seen += 1
+                assert step == 4
+            
+        if done:
+            break
+            
+    # Should see exactly 1 thermal shock at step 4
+    assert events_seen == 1, f"Expected 1 thermal shock, got {events_seen}"

utils/constants.py

@@ -47,16 +47,15 @@ DEFAULT_STATE: Dict[str, float] = {
 }
 
 # Episode controls.
-# Shorter episode length helps validation runs show signal quickly.
-DEFAULT_MAX_STEPS = 12
+# Tasks dictate default episode lengths. Max steps are defined in the specific Task subclasses.
 DEFAULT_TASK_ID = "task1"
-DEFAULT_EPISODE_ID = 0
+DEFAULT_EPISODE_ID = 143
 OBSERVATION_DECIMALS = 2
 
 # External evaluator episode policy
 # External (evaluator/public) calls will use this fixed episode id unless
 # a valid developer token is supplied via the `X-DEV-TOKEN` header.
-DEFAULT_EXTERNAL_EPISODE_ID = 1
+DEFAULT_EXTERNAL_EPISODE_ID = 143
 EXTERNAL_EPISODE_DISPLAY_WIDTH = 3
 
 # Name of the environment variable that holds the developer reset token.
@@ -110,75 +109,75 @@ INIT_T_F_GAIN = 0.30
 TASK_STARTUP_PROFILES: Dict[str, Dict[str, float]] = {
 	"task1": {
 		"aL": 0.60,
-		"sL": 0.10,
+		"sL": 0.00, # Fixed load start to match constant task target
+		"aD": 0.02,
+		"sD": 0.03,
+		"aT": 0.08,
+		"sT": 0.06,
+		"aM": 0.88,
+		"sM": 0.06,
+		"gap_scale": 0.22,
+		"d_max": 0.08,
+		"s_base": 0.05,
+		"s_gain": 0.08,
+		"f_bias": 0.20,
+		"t_task_bias": -0.08,
+		"soft_t_cap": 0.82,
+		"soft_pr_cap": 0.86,
+	},
+	"task2": {
+		"aL": 0.50, # Set to exact step 1 load sequence baseline
+		"sL": 0.00,
 		"aD": 0.04,
-		"sD": 0.05,
-		"aT": 0.12,
+		"sD": 0.04,
+		"aT": 0.14,
 		"sT": 0.10,
-		"aM": 0.82,
+		"aM": 0.78,
 		"sM": 0.08,
-		"gap_scale": 0.05,
-		"d_max": 0.15,
-		"s_base": 0.10,
-		"s_gain": 0.14,
-		"f_bias": 0.16,
-		"t_task_bias": -0.06,
-		"soft_t_cap": 0.78,
-		"soft_pr_cap": 0.84,
-	},
-	"task2": {
-		"aL": 0.50,
-		"sL": 0.12,
-		"aD": 0.05,
-		"sD": 0.06,
-		"aT": 0.18,
-		"sT": 0.14,
-		"aM": 0.70,
-		"sM": 0.10,
-		"gap_scale": 0.08,
+		"gap_scale": 0.16,
 		"d_max": 0.10,
-		"s_base": 0.00,
-		"s_gain": 0.18,
-		"f_bias": 0.18,
-		"t_task_bias": -0.02,
-		"soft_t_cap": 0.82,
+		"s_base": 0.03,
+		"s_gain": 0.12,
+		"f_bias": 0.22,
+		"t_task_bias": -0.04,
+		"soft_t_cap": 0.84,
 		"soft_pr_cap": 0.88,
 	},
 	"task3": {
 		"aL": 0.70,
-		"sL": 0.10,
-		"aD": 0.12,
-		"sD": 0.10,
-		"aT": 0.48,
-		"sT": 0.18,
-		"aM": 0.46,
-		"sM": 0.10,
-		"gap_scale": 0.07,
-		"d_max": 0.22,
-		"s_base": 0.08,
-		"s_gain": 0.30,
-		"f_bias": 0.24,
-		"t_task_bias": 0.05,
-		"soft_t_cap": 0.92,
-		"soft_pr_cap": 0.98,
+		"sL": 0.00,
+		"aD": 0.08,
+		"sD": 0.06,
+		"aT": 0.72,
+		"sT": 0.12,
+		"aM": 0.60,
+		"sM": 0.08,
+		"gap_scale": 0.05, # Reduce gap further, so power is already high and generating heat
+		"d_max": 0.18,
+		"s_base": 0.20,
+		"s_gain": 0.18,
+		"f_bias": 0.12,
+		"t_task_bias": 0.35, # Extra direct heat
+		"soft_t_cap": 1.05,  # Prevent environment from rescuing the initial state
+		"soft_pr_cap": 1.10,
 	},
 	"task4": {
 		"aL": 0.60,
-		"sL": 0.10,
-		"aD": 0.34,
-		"sD": 0.14,
-		"aT": 0.34,
-		"sT": 0.16,
-		"aM": 0.55,
-		"sM": 0.10,
-		"gap_scale": 0.06,
-		"d_max": 0.45,
-		"s_base": 0.03,
-		"s_gain": 0.22,
-		"f_bias": 0.30,
-		"t_task_bias": 0.04,
-		"soft_t_cap": 0.88,
-		"soft_pr_cap": 0.94,
+		"sL": 0.00,
+		"aD": 0.18,
+		"sD": 0.08,
+		"aT": 0.22,
+		"sT": 0.10,
+		"aM": 0.68,
+		"sM": 0.08,
+		"gap_scale": 0.14,
+		"d_max": 0.32,
+		"s_base": 0.05,
+		"s_gain": 0.14,
+		"f_bias": 0.32,
+		"t_task_bias": 0.02,
+		"soft_t_cap": 0.85,
+		"soft_pr_cap": 0.89,
 	},
 }