fix: conform to OpenEnv base interface contract

inference.py

@@ -115,7 +115,9 @@ def run_task1_pairwise(env_client) -> dict[str, Any]:
 
         from models import PairwiseAction
         action = PairwiseAction(choice=choice)
-        obs, reward, done, info = env_client.step(action)
+        obs = env_client.step(action)
+        reward = obs.reward
+        done = obs.done
 
         total_reward += reward
         steps += 1
@@ -168,7 +170,9 @@ def run_task2_likert(env_client) -> dict[str, Any]:
             harmlessness=clamp("harmlessness"),
             instruction_following=clamp("instruction_following"),
         )
-        obs, reward, done, info = env_client.step(action)
+        obs = env_client.step(action)
+        reward = obs.reward
+        done = obs.done
 
         total_reward += reward
         steps += 1
@@ -215,7 +219,9 @@ def run_task3_consistency(env_client) -> dict[str, Any]:
 
         from models import ConsistencyAction
         action = ConsistencyAction(ranking=ranking)
-        obs, reward, done, info = env_client.step(action)
+        obs = env_client.step(action)
+        reward = obs.reward
+        done = obs.done
 
         total_reward += reward
         steps += 1

server/environment.py

@@ -2,9 +2,9 @@
 PreferenceLab Core Environment.
 
 Implements the OpenEnv Environment base class with:
-    - reset()  → returns initial observation
-    - step()   → executes action, returns (observation, reward, done, info)
-    - state()  → returns episode metadata
+    - reset()  → returns initial Observation
+    - step()   → executes action, returns Observation (reward & done embedded)
+    - state    → @property returning a State object with episode metadata
 
 Three tasks:
     Task 1 (pairwise)     - Easy:   pairwise choice graded against HH-RLHF gold labels
@@ -20,6 +20,7 @@ from pathlib import Path
 from typing import Any
 
 from openenv.core.env_server import Environment
+from openenv.core.env_server.types import State
 
 from models import (
     ConsistencyAction,
@@ -258,7 +259,8 @@ class PreferenceLabEnvironment(Environment):
             timeout_s: Unused — required by base class signature.
 
         Returns:
-            Tuple of (observation, reward, done, info).
+            Observation with reward and done embedded as fields.
+            Read obs.reward and obs.done instead of unpacking a tuple.
         """
         self._step_count += 1
 
@@ -273,19 +275,19 @@ class PreferenceLabEnvironment(Environment):
             rng = random.Random(self._seed + self._step_count)
             self._current_example = self._sample_example(rng)
 
-        obs = self._build_observation(reward=reward, done=done, info=info)
-        return obs, reward, done, info
-
-    def state(self) -> dict[str, Any]:
-        """Return current episode metadata."""
-        return {
-            "episode_id": self._episode_id,
-            "step_count": self._step_count,
-            "task_type": self._task_type,
-            "cumulative_reward": round(self._cumulative_reward, 4),
-            "max_steps": MAX_STEPS_PER_EPISODE,
-            "seed": self._seed,
-        }
+        return self._build_observation(reward=reward, done=done, info=info)
+
+    @property
+    def state(self) -> State:
+        """Return current episode metadata as an openenv State object."""
+        return State(
+            episode_id=self._episode_id,
+            step_count=self._step_count,
+            task_type=self._task_type,
+            cumulative_reward=round(self._cumulative_reward, 4),
+            max_steps=MAX_STEPS_PER_EPISODE,
+            seed=self._seed,
+        )
 
     # ── Internal helpers ──────────────────────────────────────
 

test_api.py

@@ -3,20 +3,24 @@ from models import PairwiseAction, LikertAction, ConsistencyAction
 
 env = PreferenceLabEnvironment()
 
-# Task 1
+# Task 1 – Pairwise
 obs = env.reset(seed=42, task_type='pairwise')
 print('TASK 1 prompt:', obs.prompt[:60])
-obs, r, done, _ = env.step(PairwiseAction(choice='A'))
-print('Reward:', r)
+obs = env.step(PairwiseAction(choice='A'))
+print('Reward:', obs.reward, '| Done:', obs.done)
 
-# Task 2  
+# Task 2 – Likert
 obs = env.reset(seed=42, task_type='likert')
-print('TASK 2 response:', obs.response[:60])
-obs, r, done, _ = env.step(LikertAction(helpfulness=5, honesty=5, harmlessness=5, instruction_following=5))
-print('Reward:', r)
+print('\nTASK 2 response:', obs.response[:60])
+obs = env.step(LikertAction(helpfulness=5, honesty=5, harmlessness=5, instruction_following=5))
+print('Reward:', obs.reward, '| Done:', obs.done)
 
-# Task 3
+# Task 3 – Consistency
 obs = env.reset(seed=42, task_type='consistency')
-print('TASK 3 prompt:', obs.prompt[:60])
-obs, r, done, _ = env.step(ConsistencyAction(ranking=['C','A','B','D']))
-print('Reward:', r)
+print('\nTASK 3 prompt:', obs.prompt[:60])
+obs = env.step(ConsistencyAction(ranking=['C', 'A', 'B', 'D']))
+print('Reward:', obs.reward, '| Done:', obs.done)
+
+# State (now a property, not a method call)
+state = env.state
+print('\nState:', state.model_dump())

tests/test_environment.py

@@ -182,42 +182,43 @@ class TestPreferenceLabEnvironment:
     def test_step_pairwise(self):
         self.env.reset(task_type="pairwise")
         action = PairwiseAction(choice="A")
-        obs, reward, done, info = self.env.step(action)
+        obs = self.env.step(action)
         assert isinstance(obs, PairwiseObservation)
-        assert 0.0 <= reward <= 1.0
-        assert isinstance(done, bool)
+        assert 0.0 <= obs.reward <= 1.0
+        assert isinstance(obs.done, bool)
 
     def test_step_likert(self):
         self.env.reset(task_type="likert")
         action = LikertAction(helpfulness=4, honesty=4, harmlessness=5, instruction_following=4)
-        obs, reward, done, info = self.env.step(action)
+        obs = self.env.step(action)
         assert isinstance(obs, LikertObservation)
-        assert 0.0 <= reward <= 1.0
+        assert 0.0 <= obs.reward <= 1.0
 
     def test_step_consistency(self):
         self.env.reset(task_type="consistency")
         action = ConsistencyAction(ranking=["A", "B", "C", "D"])
-        obs, reward, done, info = self.env.step(action)
+        obs = self.env.step(action)
         assert isinstance(obs, ConsistencyObservation)
-        assert 0.0 <= reward <= 1.0
+        assert 0.0 <= obs.reward <= 1.0
 
     def test_episode_terminates_after_max_steps(self):
         self.env.reset(task_type="pairwise")
         done = False
         steps = 0
         while not done:
-            _, _, done, _ = self.env.step(PairwiseAction(choice="A"))
+            obs = self.env.step(PairwiseAction(choice="A"))
+            done = obs.done
             steps += 1
             assert steps <= 10, "Episode ran too long!"
         assert done is True
 
     def test_state_returns_metadata(self):
         self.env.reset(seed=42, task_type="pairwise")
-        state = self.env.state()
-        assert "episode_id" in state
-        assert "step_count" in state
-        assert "task_type" in state
-        assert state["seed"] == 42
+        state = self.env.state
+        assert "episode_id" in state.model_dump()
+        assert "step_count" in state.model_dump()
+        assert "task_type" in state.model_dump()
+        assert state.seed == 42
 
     def test_reproducible_with_seed(self):
         obs1 = self.env.reset(seed=123, task_type="pairwise")
@@ -230,13 +231,11 @@ class TestPreferenceLabEnvironment:
         rewards = set()
         for _ in range(5):
             self.env.reset(task_type="pairwise")
-            action_a = PairwiseAction(choice="A")
-            _, r1, _, _ = self.env.step(action_a)
+            obs_a = self.env.step(PairwiseAction(choice="A"))
             self.env.reset(task_type="pairwise")
-            action_b = PairwiseAction(choice="B")
-            _, r2, _, _ = self.env.step(action_b)
-            rewards.add(r1)
-            rewards.add(r2)
+            obs_b = self.env.step(PairwiseAction(choice="B"))
+            rewards.add(obs_a.reward)
+            rewards.add(obs_b.reward)
         assert len(rewards) > 1, "Grader always returns the same score — DISQUALIFICATION!"
 
     def test_all_three_tasks_run(self):