Upload folder using huggingface_hub

models.py

@@ -31,9 +31,17 @@ class RewardDetails(BaseModel):
     test_reward: float = Field(default=0.0, description="Reward from passing tests")
     quality_bonus: float = Field(default=0.0, description="Bonus for code quality improvements")
     correctness_bonus: float = Field(default=0.0, description="Bonus for full correctness")
+    progress_delta: float = Field(default=0.0, description="Reward from score improvement")
+    stagnation_penalty: float = Field(default=0.0, description="Penalty for code not changing")
+    regression_penalty: float = Field(default=0.0, description="Penalty for score decline")
     invalid_action_penalty: float = Field(default=0.0, description="Penalty for invalid actions")
     timeout_penalty: float = Field(default=0.0, description="Penalty for timeouts")
     reason: str = Field(..., description="Explanation of reward")
+    
+    # Debug info
+    prev_score: float = Field(default=0.0, description="Score before this step")
+    curr_score: float = Field(default=0.0, description="Score after this step")
+    code_changed: bool = Field(default=False, description="Whether code was modified")
 
 
 class PythonCodeReviewAction(Action):
@@ -43,26 +51,26 @@ class PythonCodeReviewAction(Action):
     code: Optional[str] = Field(default=None, description="New code for edit_code actions")
 
 
-class PythonCodeReviewObservation(Observation):
-    """Observation returned by reset() and step()."""
-
-    task_id: str = Field(..., description="Current task identifier")
-    title: str = Field(default="", description="Human-readable task title")
-    difficulty: Difficulty = Field(..., description="Task difficulty level")
-    task_kind: Optional[TaskKind] = Field(default=None, description="Task type")
-    task_description: str = Field(..., description="Detailed task description")
-    current_code: str = Field(..., description="Current code state")
-    errors: str = Field(..., description="Syntax/compilation errors, if any")
-    test_results: str = Field(..., description="Results from test execution")
-    visible_tests: List[str] = Field(default_factory=list, description="Public test cases")
-    history: List[HistoryEntry] = Field(default_factory=list, description="Action history")
-    attempts_remaining: int = Field(..., ge=0, description="Actions left in episode")
-    last_action_status: str = Field(default="", description="Outcome message from the last action")
-    score: float = Field(..., ge=0.0, le=1.0, description="Current episode score")
-    reward_details: RewardDetails = Field(
-        default_factory=lambda: RewardDetails(value=0.0, reason="Reset"),
-        description="Detailed reward breakdown for the last action",
-    )
+class PythonCodeReviewObservation(Observation):
+    """Observation returned by reset() and step()."""
+
+    task_id: str = Field(..., description="Current task identifier")
+    title: str = Field(default="", description="Human-readable task title")
+    difficulty: Difficulty = Field(..., description="Task difficulty level")
+    task_kind: Optional[TaskKind] = Field(default=None, description="Task type")
+    task_description: str = Field(..., description="Detailed task description")
+    current_code: str = Field(..., description="Current code state")
+    errors: str = Field(..., description="Syntax/compilation errors, if any")
+    test_results: str = Field(..., description="Results from test execution")
+    visible_tests: List[str] = Field(default_factory=list, description="Public test cases")
+    history: List[HistoryEntry] = Field(default_factory=list, description="Action history")
+    attempts_remaining: int = Field(..., ge=0, description="Actions left in episode")
+    last_action_status: str = Field(default="", description="Outcome message from the last action")
+    score: float = Field(..., ge=0.0, le=1.0, description="Current episode score")
+    reward_details: RewardDetails = Field(
+        default_factory=lambda: RewardDetails(value=0.0, reason="Reset"),
+        description="Detailed reward breakdown for the last action",
+    )
 
 
 class PythonCodeReviewState(State):
@@ -112,4 +120,4 @@ class HealthResponse(BaseModel):
 
     status: Literal["ok"] = "ok"
     environment: str = "python_code_review_env"
-    task_count: int = Field(default=0, ge=0)
+    task_count: int = Field(default=0, ge=0)

server/env.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import sys
 from typing import List, Optional
 from uuid import uuid4
 
@@ -20,13 +21,16 @@ from models import (
 from tasks import TaskSpec, get_task, list_task_descriptors, list_task_summaries, task_ids
 
 
-# Reward shaping constants
-INVALID_ACTION_PENALTY = 0.1
-QUALITY_BONUS_SCALE = 0.15
-ANALYZE_FAILURE_PENALTY = 0.05
-RUN_FAILURE_PENALTY = 0.05
-TIMEOUT_PENALTY = 0.1
-SUBMIT_BASE_SCALE = 0.1
+# Reward shaping constants (balanced for meaningful variation)
+SYNTAX_FIX_BONUS = 0.35  # One-time for fixing syntax
+TEST_PASS_REWARD_SCALE = 0.30  # Per test improvement
+QUALITY_BONUS_SCALE = 0.15  # Code quality improvement
+PROGRESS_SCALE = 0.25  # Score improvement
+COMPLETION_BONUS = 0.50  # Full correctness (one-time)
+INVALID_ACTION_PENALTY = 0.15
+STAGNATION_PENALTY = 0.10  # If code unchanged but action taken
+REGRESSION_PENALTY_SCALE = 0.20  # Per 0.1 score decline
+TIMEOUT_PENALTY = 0.15
 
 
 class PythonCodeReviewEnvironment(
@@ -36,17 +40,22 @@ class PythonCodeReviewEnvironment(
 
     SUPPORTS_CONCURRENT_SESSIONS = True
 
-    def __init__(self) -> None:
-        super().__init__()
-        self._task_order = list(task_ids())
-        self._task_cursor = -1
-        self._task: Optional[TaskSpec] = None
-        self._state = PythonCodeReviewState(episode_id=str(uuid4()))
-        self._done = False
-        self._last_status = "Call reset() to start."
-        self._last_reward = RewardDetails(value=0.0, reason="Environment initialized.")
-        self._best_visible_test_fraction = 0.0
-        self._best_quality_score = 0.0
+    def __init__(self, verbose: bool = True) -> None:
+        super().__init__()
+        self._task_order = list(task_ids())
+        self._task_cursor = -1
+        self._task: Optional[TaskSpec] = None
+        self._state = PythonCodeReviewState(episode_id=str(uuid4()))
+        self._done = False
+        self._last_status = "Call reset() to start."
+        self._last_reward = RewardDetails(value=0.0, reason="Environment initialized.")
+        self._verbose = verbose
+        
+        # Progress tracking
+        self._previous_score = 0.0
+        self._previous_code = ""
+        self._best_visible_test_fraction = 0.0
+        self._best_quality_score = 0.0
         self._full_correctness_awarded = False
         self._syntax_reward_awarded = False
 
@@ -69,14 +78,16 @@ class PythonCodeReviewEnvironment(
             self._task_cursor = (self._task_cursor + 1) % len(self._task_order)
             self._task = get_task(self._task_order[self._task_cursor])
 
-        # Reset episode state
+        # Reset episode state and tracking
         self._done = False
+        self._previous_score = 0.0
+        self._previous_code = self._task.starter_code
         self._best_visible_test_fraction = 0.0
         self._best_quality_score = 0.0
         self._full_correctness_awarded = False
         self._syntax_reward_awarded = False
         self._last_status = "Inspect the code, edit it, run tests, then submit."
-        self._last_reward = RewardDetails(value=0.0, reason="Episode reset.")
+        self._last_reward = RewardDetails(value=0.0, reason="Episode reset.", prev_score=0.0, curr_score=0.0)
         
         self._state = PythonCodeReviewState(
             episode_id=episode_id or str(uuid4()),
@@ -93,6 +104,11 @@ class PythonCodeReviewEnvironment(
             done=False,
         )
         
+        if self._verbose:
+            print(f"\n{'='*70}")
+            print(f"RESET: Task {self._task.task_id} ({self._task.difficulty})")
+            print(f"{'='*70}")
+        
         return self._build_observation()
 
     def step(
@@ -148,6 +164,10 @@ class PythonCodeReviewEnvironment(
             self._finalize_episode(auto_submit=True)
             self._state.done = True
 
+        # Debug logging
+        if self._verbose:
+            self._log_debug_step(reward)
+
         return self._build_observation()
 
     @property
@@ -155,13 +175,9 @@ class PythonCodeReviewEnvironment(
         """Return the current environment state."""
         return self._state.model_copy(deep=True)
 
-    def list_task_summaries(self) -> List[object]:
-        """Return public task metadata."""
-        return list_task_summaries()
-
-    def list_tasks(self) -> List[object]:
-        """Return all public task descriptors."""
-        return list_task_descriptors()
+    def list_task_summaries(self) -> List[object]:
+        """Return public task metadata."""
+        return list_task_summaries()
 
     def get_task(self, task_id: str) -> object:
         """Return a single task descriptor."""
@@ -175,60 +191,79 @@ class PythonCodeReviewEnvironment(
         """Expose deterministic grading outside of an active episode."""
         return grade_task(code, get_task(task_id), include_hidden=True)
 
-    def _build_observation(self) -> PythonCodeReviewObservation:
-        """Build current observation from state."""
-        return PythonCodeReviewObservation(
-            task_id=self._state.task_id or "",
-            title=self._task.title if self._task else "",
-            difficulty=self._state.difficulty or "easy",
-            task_kind=self._state.task_kind,
-            task_description=self._task.task_description if self._task else "",
-            current_code=self._state.current_code,
-            errors=self._state.errors,
-            test_results=self._state.test_results,
-            visible_tests=self._task.visible_tests if self._task else [],
-            history=self._state.history,
-            attempts_remaining=self._state.attempts_remaining,
-            last_action_status=self._last_status,
-            score=self._state.score,
-            reward=self._last_reward.value,
-            reward_details=self._last_reward,
-            done=self._done,
-            metadata={
-                "episode_id": self._state.episode_id,
-                "step_count": self._state.step_count,
-                "task_kind": self._state.task_kind,
-            },
-        )
-
-    def _handle_analyze(self) -> tuple[RewardDetails, str]:
-        """Analyze code for errors and test status."""
-        if self._task is None:
-            return RewardDetails(value=0.0, reason="Invalid state"), "Error: task not loaded"
-        
-        grade = grade_task(self._state.current_code, self._task, include_hidden=False)
-        error = grade.details.get("compile_error", "")
+    def _build_observation(self) -> PythonCodeReviewObservation:
+        """Build current observation from state."""
+        return PythonCodeReviewObservation(
+            task_id=self._state.task_id or "",
+            title=self._task.title if self._task else "",
+            difficulty=self._state.difficulty or "easy",
+            task_kind=self._state.task_kind,
+            task_description=self._task.task_description if self._task else "",
+            current_code=self._state.current_code,
+            errors=self._state.errors,
+            test_results=self._state.test_results,
+            visible_tests=self._task.visible_tests if self._task else [],
+            history=self._state.history,
+            attempts_remaining=self._state.attempts_remaining,
+            last_action_status=self._last_status,
+            score=self._state.score,
+            reward_details=self._last_reward,
+        )
+
+    def _handle_analyze(self) -> tuple[RewardDetails, str]:
+        """Analyze code for errors and test status."""
+        if self._task is None:
+            return RewardDetails(value=0.0, reason="Invalid state"), "Error: task not loaded"
         
+        # Grade current code
+        curr_grade = grade_task(self._state.current_code, self._task, include_hidden=False)
+        curr_score = curr_grade.score
+        error = curr_grade.details.get("compile_error", "")
+        
+        # Compute reward components
+        reward_components = self._compute_reward_components(
+            curr_score=curr_score,
+            prev_score=self._previous_score,
+            curr_grade=curr_grade,
+            code_changed=False,  # analyze doesn't change code
+        )
+        
+        # Status message
         if error:
             self._state.errors = error
             self._state.test_results = "Compilation failed. Fix syntax first."
             summary = f"Syntax error detected: {error}"
-        else:
-            self._state.errors = ""
-            if self._task.task_kind == "syntax_fix":
-                self._state.test_results = "Code compiles successfully."
-                summary = "Code compiles. Ready to submit."
-            else:
-                visible_total = len(self._task.visible_tests)
-                visible_passed = grade.tests_passed
-                self._state.test_results = f"Test run: {visible_passed}/{visible_total} passing."
-                summary = self._state.test_results
-
-        reward_value = round((grade.score * 0.05) - self._grade_penalty(grade, failure_penalty=ANALYZE_FAILURE_PENALTY), 6)
-        reward = RewardDetails(value=reward_value, reason=summary)
-        self._append_history("analyze_code", summary, reward.value)
-        self._sync_score(include_hidden=False)
-        return reward, summary
+        else:
+            self._state.errors = ""
+            if self._task.task_kind == "syntax_fix":
+                self._state.test_results = "Code compiles successfully."
+                summary = "Code compiles. Ready to submit."
+            else:
+                visible_total = len(self._task.visible_tests)
+                visible_passed = curr_grade.tests_passed
+                self._state.test_results = f"Test run: {visible_passed}/{visible_total} passing."
+                summary = self._state.test_results
+
+        reward = RewardDetails(
+            value=reward_components["total"],
+            progress_delta=reward_components["progress"],
+            syntax_reward=reward_components["syntax"],
+            test_reward=reward_components["test"],
+            quality_bonus=reward_components["quality"],
+            stagnation_penalty=reward_components["stagnation"],
+            regression_penalty=reward_components["regression"],
+            timeout_penalty=TIMEOUT_PENALTY if curr_grade.timed_out else 0.0,
+            reason=summary,
+            prev_score=round(self._previous_score, 4),
+            curr_score=round(curr_score, 4),
+            code_changed=False,
+        )
+        
+        # Update state
+        self._state.score = curr_score
+        self._state.errors = curr_grade.details.get("compile_error", "")
+        self._append_history("analyze_code", summary, reward.value)
+        return reward, summary
 
     def _handle_edit(self, action: PythonCodeReviewAction) -> tuple[RewardDetails, str]:
         """Edit the code and compute reward for progress."""
@@ -246,109 +281,218 @@ class PythonCodeReviewEnvironment(
             self._append_history("edit_code", status, reward.value)
             return reward, status
 
-        # Grade before and after
-        previous_grade = grade_task(self._state.current_code, self._task, include_hidden=False)
-        new_grade = grade_task(code, self._task, include_hidden=False)
-        self._state.current_code = code
-
+        # Detect code change
+        code_changed = (code != self._previous_code)
+        
+        # Grade before and after
+        curr_grade = grade_task(code, self._task, include_hidden=False)
+        curr_score = curr_grade.score
+        
+        # Compute reward components
+        reward_components = self._compute_reward_components(
+            curr_score=curr_score,
+            prev_score=self._previous_score,
+            curr_grade=curr_grade,
+            code_changed=code_changed,
+            prev_grade_score=grade_task(self._previous_code, self._task, include_hidden=False).syntax_score,
+        )
+        
         # Update state
-        self._state.errors = new_grade.details.get("compile_error", "")
-        self._state.test_results = self._format_test_results(new_grade)
-
-        # Compute reward with shaping
-        syntax_reward = 0.0
-        if previous_grade.syntax_score < 1.0 and new_grade.syntax_score == 1.0:
-            syntax_reward = 0.2
-            self._syntax_reward_awarded = True
-
-        quality_delta = new_grade.quality_score - previous_grade.quality_score
-        quality_bonus = max(min(quality_delta * QUALITY_BONUS_SCALE, 0.1), -0.1)
-        if new_grade.quality_score > self._best_quality_score:
-            self._best_quality_score = new_grade.quality_score
-
-        progress_reward = 0.2 * (new_grade.score - previous_grade.score)
-        if new_grade.tests_total > 0:
-            current_test_fraction = new_grade.tests_passed / new_grade.tests_total
-            self._best_visible_test_fraction = max(self._best_visible_test_fraction, current_test_fraction)
-
-        penalty = self._grade_penalty(new_grade)
-        reward_value = round(progress_reward + syntax_reward + quality_bonus - penalty, 6)
-        
-        status = "Code updated."
-        if self._state.errors:
-            status = f"Code updated with syntax issues: {self._state.errors}"
-        elif new_grade.tests_total > 0:
+        self._state.current_code = code
+        self._previous_code = code
+        self._previous_score = curr_score
+        self._state.errors = curr_grade.details.get("compile_error", "")
+        self._state.test_results = self._format_test_results(curr_grade)
+        self._state.score = curr_score
+
+        status = "Code updated."
+        if self._state.errors:
+            status = f"Code updated, but syntax issues remain: {self._state.errors}"
+        elif curr_grade.tests_total > 0:
             status = self._state.test_results
 
-        reward = RewardDetails(
-            value=reward_value,
-            syntax_reward=syntax_reward,
-            quality_bonus=round(quality_bonus, 6),
-            test_reward=round(progress_reward, 6),
-            timeout_penalty=TIMEOUT_PENALTY if new_grade.timed_out else 0.0,
-            reason=status,
-        )
-        self._append_history("edit_code", status, reward_value)
-        self._sync_score(include_hidden=False)
-        return reward, status
+        reward = RewardDetails(
+            value=reward_components["total"],
+            progress_delta=reward_components["progress"],
+            syntax_reward=reward_components["syntax"],
+            test_reward=reward_components["test"],
+            quality_bonus=reward_components["quality"],
+            stagnation_penalty=reward_components["stagnation"],
+            regression_penalty=reward_components["regression"],
+            timeout_penalty=TIMEOUT_PENALTY if curr_grade.timed_out else 0.0,
+            reason=status,
+            prev_score=round(self._previous_score - curr_score + self._previous_score, 4) if curr_score != self._previous_score else round(self._previous_score, 4),
+            curr_score=round(curr_score, 4),
+            code_changed=code_changed,
+        )
+        
+        self._append_history("edit_code", status, reward.value)
+        return reward, status
 
     def _handle_run_tests(self) -> tuple[RewardDetails, str]:
         """Run tests and provide feedback."""
         if self._task is None:
             return RewardDetails(value=0.0, reason="Invalid state"), "Error: task not loaded"
         
-        grade = grade_task(self._state.current_code, self._task, include_hidden=False)
-        self._state.errors = grade.details.get("compile_error", "")
-        self._state.test_results = self._format_test_results(grade)
-
-        previous_score = self._state.score
-        progress_reward = 0.2 * (grade.score - previous_score)
-        completion_bonus = 0.05 if grade.tests_total > 0 and grade.tests_passed == grade.tests_total else 0.0
-        penalty = self._grade_penalty(grade, failure_penalty=RUN_FAILURE_PENALTY)
-        reward_value = round(progress_reward + completion_bonus - penalty, 6)
-        if grade.tests_total > 0:
-            current_fraction = grade.tests_passed / grade.tests_total
-            self._best_visible_test_fraction = max(self._best_visible_test_fraction, current_fraction)
-
-        status = self._state.test_results if not self._state.errors else self._state.errors
-        reward = RewardDetails(
-            value=reward_value,
-            test_reward=round(progress_reward + completion_bonus, 6),
-            timeout_penalty=TIMEOUT_PENALTY if grade.timed_out else 0.0,
-            reason=status,
-        )
-        self._append_history("run_tests", status, reward.value)
-        self._sync_score(include_hidden=False)
-        return reward, status
+        curr_grade = grade_task(self._state.current_code, self._task, include_hidden=False)
+        curr_score = curr_grade.score
+        
+        # Compute reward components
+        reward_components = self._compute_reward_components(
+            curr_score=curr_score,
+            prev_score=self._previous_score,
+            curr_grade=curr_grade,
+            code_changed=True,  # Consider any test run as "changed behavior"
+        )
+        
+        # Update state
+        self._state.errors = curr_grade.details.get("compile_error", "")
+        self._state.test_results = self._format_test_results(curr_grade)
+        self._state.score = curr_score
+        self._previous_score = curr_score
+
+        # Add completion bonus if all visible tests pass
+        completion_bonus = 0.0
+        if curr_grade.tests_total > 0 and curr_grade.tests_passed == curr_grade.tests_total:
+            completion_bonus = 0.20
+
+        status = self._state.test_results if not self._state.errors else self._state.errors
+        reward = RewardDetails(
+            value=reward_components["total"] + completion_bonus,
+            progress_delta=reward_components["progress"],
+            test_reward=reward_components["test"] + completion_bonus,
+            syntax_reward=reward_components["syntax"],
+            quality_bonus=reward_components["quality"],
+            stagnation_penalty=reward_components["stagnation"],
+            regression_penalty=reward_components["regression"],
+            timeout_penalty=TIMEOUT_PENALTY if curr_grade.timed_out else 0.0,
+            reason=status,
+            prev_score=round(self._previous_score - curr_score + self._previous_score, 4) if curr_score != self._previous_score else round(self._previous_score, 4),
+            curr_score=round(curr_score, 4),
+            code_changed=True,
+        )
+        
+        self._append_history("run_tests", status, reward.value)
+        return reward, status
 
     def _handle_submit(self) -> tuple[RewardDetails, str]:
         """Submit solution and finalize episode."""
         if self._task is None:
             return RewardDetails(value=0.0, reason="Invalid state"), "Error: task not loaded"
         
-        grade = grade_task(self._state.current_code, self._task, include_hidden=True)
-        self._state.errors = grade.details.get("compile_error", "")
-        self._state.test_results = self._format_test_results(grade)
-
-        # Compute final reward bonuses
-        correctness_bonus = 0.0
-        if grade.score >= 0.999999 and not self._full_correctness_awarded:
-            correctness_bonus = 0.5
-            self._full_correctness_awarded = True
-
-        penalty = self._grade_penalty(grade, failure_penalty=RUN_FAILURE_PENALTY)
-        reward_value = round((grade.score * SUBMIT_BASE_SCALE) + correctness_bonus - penalty, 6)
-        self._finalize_episode(auto_submit=False, grade=grade)
-        status = f"Solution submitted. Final score: {grade.score:.3f}"
-        
-        reward = RewardDetails(
-            value=reward_value,
-            correctness_bonus=correctness_bonus,
-            timeout_penalty=TIMEOUT_PENALTY if grade.timed_out else 0.0,
-            reason=status,
-        )
-        self._append_history("submit_solution", status, reward_value)
-        return reward, status
+        curr_grade = grade_task(self._state.current_code, self._task, include_hidden=True)
+        curr_score = curr_grade.score
+        
+        # Compute reward components
+        reward_components = self._compute_reward_components(
+            curr_score=curr_score,
+            prev_score=self._previous_score,
+            curr_grade=curr_grade,
+            code_changed=False,  # Submit doesn't change code
+        )
+        
+        # Bonus for full correctness (one-time only)
+        correctness_bonus = 0.0
+        if curr_score >= 0.9999 and not self._full_correctness_awarded:
+            correctness_bonus = COMPLETION_BONUS
+            self._full_correctness_awarded = True
+        
+        # Update state
+        self._state.errors = curr_grade.details.get("compile_error", "")
+        self._state.test_results = self._format_test_results(curr_grade)
+        self._state.score = curr_score
+        self._finalize_episode(auto_submit=False, grade=curr_grade)
+        
+        reward_value = max(-1.0, min(1.0, reward_components["total"] + correctness_bonus))
+        status = f"Solution submitted. Final score: {curr_score:.3f}"
+        
+        reward = RewardDetails(
+            value=reward_value,
+            progress_delta=reward_components["progress"],
+            correctness_bonus=correctness_bonus,
+            syntax_reward=reward_components["syntax"],
+            test_reward=reward_components["test"],
+            quality_bonus=reward_components["quality"],
+            stagnation_penalty=reward_components["stagnation"],
+            regression_penalty=reward_components["regression"],
+            timeout_penalty=TIMEOUT_PENALTY if curr_grade.timed_out else 0.0,
+            reason=status,
+            prev_score=round(self._previous_score, 4),
+            curr_score=round(curr_score, 4),
+            code_changed=False,
+        )
+        
+        self._append_history("submit_solution", status, reward_value)
+        return reward, status
+
+    def _compute_reward_components(
+        self,
+        curr_score: float,
+        prev_score: float,
+        curr_grade: TaskGrade,
+        code_changed: bool,
+        prev_grade_score: float = 0.0,
+    ) -> dict:
+        """Compute all reward components and return as dict."""
+        components = {
+            "progress": 0.0,
+            "syntax": 0.0,
+            "test": 0.0,
+            "quality": 0.0,
+            "stagnation": 0.0,
+            "regression": 0.0,
+            "total": 0.0,
+        }
+        
+        # 1. Progress reward: score improvement
+        score_delta = curr_score - prev_score
+        if score_delta > 0:
+            components["progress"] = min(PROGRESS_SCALE * score_delta, 0.25)
+        
+        # 2. Syntax reward: one-time bonus for first time compiling
+        if not self._syntax_reward_awarded and curr_grade.syntax_score >= 0.99:
+            if prev_grade_score < 0.99:
+                components["syntax"] = SYNTAX_FIX_BONUS
+                self._syntax_reward_awarded = True
+        
+        # 3. Test reward: improvement in test pass rate
+        if curr_grade.tests_total > 0:
+            curr_test_frac = curr_grade.tests_passed / curr_grade.tests_total
+            test_delta = curr_test_frac - self._best_visible_test_fraction
+            if test_delta > 0:
+                components["test"] = min(TEST_PASS_REWARD_SCALE * test_delta, 0.20)
+                self._best_visible_test_fraction = max(self._best_visible_test_fraction, curr_test_frac)
+        
+        # 4. Quality reward: code quality improvement
+        quality_delta = curr_grade.quality_score - self._best_quality_score
+        if quality_delta > 0:
+            components["quality"] = min(QUALITY_BONUS_SCALE * quality_delta, 0.15)
+            self._best_quality_score = max(self._best_quality_score, curr_grade.quality_score)
+        
+        # 5. Stagnation penalty: code not changed (except analyze_code)
+        if not code_changed and not (curr_grade.details.get("compile_error") == ""):
+            components["stagnation"] = -STAGNATION_PENALTY
+        
+        # 6. Regression penalty: score decreased
+        if score_delta < 0:
+            components["regression"] = REGRESSION_PENALTY_SCALE * abs(score_delta)
+        
+        # 7. Timeout penalty
+        if curr_grade.timed_out:
+            components["regression"] = -TIMEOUT_PENALTY
+        
+        # Compute total and clamp to [-1.0, 1.0]
+        total = (
+            components["progress"]
+            + components["syntax"]
+            + components["test"]
+            + components["quality"]
+            - components["stagnation"]
+            - components["regression"]
+        )
+        components["total"] = max(-1.0, min(1.0, round(total, 6)))
+        
+        return components
 
     def _finalize_episode(self, auto_submit: bool, grade: Optional[TaskGrade] = None) -> None:
         """Mark episode as done and set final score."""
@@ -356,24 +500,10 @@ class PythonCodeReviewEnvironment(
             if self._task is None:
                 return
             grade = grade_task(self._state.current_code, self._task, include_hidden=True)
-            self._state.errors = grade.details.get("compile_error", "")
-            self._state.test_results = self._format_test_results(grade)
         
         self._state.score = grade.score
         self._done = True
         self._state.done = True
-        
-        if auto_submit:
-            self._last_status = f"Step budget exhausted. Final score: {grade.score:.3f}"
-
-    def _sync_score(self, include_hidden: bool) -> None:
-        """Update visible score based on current code."""
-        if self._task is None:
-            return
-        grade = grade_task(self._state.current_code, self._task, include_hidden=include_hidden)
-        # For visible runs, use a soft score; for hidden, it will be finalized on submit
-        if not include_hidden:
-            self._state.score = grade.score
 
     def _format_test_results(self, grade: TaskGrade) -> str:
         """Format test results for display."""
@@ -383,28 +513,42 @@ class PythonCodeReviewEnvironment(
             return "Test execution timed out."
         return f"Tests: {grade.tests_passed}/{grade.tests_total} passing"
 
-    def _append_history(self, action_type: str, status: str, reward: float) -> None:
-        """Append action to history."""
-        entry = HistoryEntry(
-            step=self._state.step_count,
-            action_type=action_type,
-            status=status,
-            reward=reward,
-        )
-        self._state.history.append(entry)
-
-    def _grade_penalty(self, grade: TaskGrade, failure_penalty: float = RUN_FAILURE_PENALTY) -> float:
-        """Return a negative signal when the action leads to an obviously bad result."""
-        penalty = 0.0
-        if grade.details.get("compile_error"):
-            penalty += failure_penalty + grade.score
-        if grade.timed_out:
-            penalty += TIMEOUT_PENALTY
-        if grade.tests_total > 0 and grade.tests_passed == 0:
-            penalty += failure_penalty
-        return round(penalty, 6)
-
-
-# Backwards-compatible aliases used elsewhere in the repo.
-PythonEnvironment = PythonCodeReviewEnvironment
-CodeReviewEnvironment = PythonCodeReviewEnvironment
+    def _append_history(self, action_type: str, status: str, reward: float) -> None:
+        """Append action to history."""
+        entry = HistoryEntry(
+            step=self._state.step_count,
+            action_type=action_type,
+            status=status,
+            reward=reward,
+        )
+        self._state.history.append(entry)
+
+    def _log_debug_step(self, reward: RewardDetails) -> None:
+        """Log step details for debugging."""
+        print(
+            f"\nStep {self._state.step_count:2d} | "
+            f"Score: {reward.curr_score:.3f} | "
+            f"Delta: {(reward.curr_score - reward.prev_score):+.3f} | "
+            f"Reward: {reward.value:+.4f} | "
+            f"Changed: {reward.code_changed}"
+        )
+        
+        # Log components if non-zero
+        components = [
+            ("Progress", reward.progress_delta),
+            ("Syntax", reward.syntax_reward),
+            ("Test", reward.test_reward),
+            ("Quality", reward.quality_bonus),
+            ("Stagnation", -reward.stagnation_penalty),
+            ("Regression", -reward.regression_penalty),
+        ]
+        
+        non_zero = [f"{name}={val:+.3f}" for name, val in components if abs(val) > 0.001]
+        if non_zero:
+            print(f"         | {' | '.join(non_zero)}")
+        print(f"         | Reason: {reward.reason}")
+
+
+# Backwards-compatible aliases used elsewhere in the repo.
+PythonEnvironment = PythonCodeReviewEnvironment
+CodeReviewEnvironment = PythonCodeReviewEnvironment