File size: 3,723 Bytes
961dad9 14b6f9b 57dcac8 14b6f9b 57dcac8 14b6f9b 961dad9 57dcac8 961dad9 57dcac8 14b6f9b 57dcac8 14b6f9b 57dcac8 961dad9 57dcac8 14b6f9b 57dcac8 14b6f9b 57dcac8 961dad9 57dcac8 14b6f9b 961dad9 14b6f9b | 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 | from env.models import Observation, Action, Reward
from env.tasks import TASKS
from env.grader import grade_detect, grade_risk, grade_fix
class CodeGuardianEnv:
def __init__(self):
self.tasks = TASKS
self.current_task_index = 0
self.current_step = 0 # 0 = detect, 1 = risk, 2 = fix
self.partial_score = 0
self.done = False
def reset(self):
self.current_task_index = 0
self.current_step = 0
self.done = False
self.partial_score = 0
return self._get_obs()
def _get_obs(self):
task = self.tasks[self.current_task_index]
return Observation(
code=task["code"],
intent=task["intent"],
constraints=task["constraints"],
task_id=task["task_id"],
difficulty=task["difficulty"],
step_index=self.current_step
)
def step(self, action: Action):
"""Execute one step in the environment with structured reward tracking"""
# Early guard to prevent IndexError when environment is done
if self.done:
reward_obj = Reward(
value=0.0,
feedback="Environment is complete. Reset to continue.",
partial_scores={},
breakdown={}
)
return None, reward_obj, self.done
task = self.tasks[self.current_task_index]
weights = [0.3, 0.2, 0.5] # detect, risk, fix weights
step_types = ['detect', 'risk', 'fix']
# Check if action order is valid
if action.action_type != step_types[self.current_step]:
score = 0.0
feedback = f"Invalid action order: expected {step_types[self.current_step]}, got {action.action_type}"
partial_scores = {}
else:
# Use deterministic graders
original_code = task["code"]
if action.action_type == "detect":
score = grade_detect(action.detected_issues or [], task["expected_issues"])
elif action.action_type == "risk":
score = grade_risk(action.risk_level or "", task["expected_risk"])
elif action.action_type == "fix":
score = grade_fix(action.suggested_fix or "", task["canonical_fix"], original_code)
else:
score = 0.0
# Weighted reward
feedback = str(score)
partial_scores = {action.action_type: score}
# Calculate weighted reward and update partial score
reward_value = score * weights[self.current_step] if partial_scores else -0.3
self.partial_score += reward_value
# Create structured reward with breakdown
reward_obj = Reward(
value=reward_value,
feedback=feedback,
partial_scores=partial_scores,
breakdown=partial_scores.copy() # Track detailed breakdown
)
# Advance step_index % 3, task if done
if partial_scores: # Only advance if action was valid
self.current_step = (self.current_step + 1) % 3
if self.current_step == 0: # Completed full cycle
self.current_task_index += 1
if self.current_task_index >= len(self.tasks):
self.done = True
next_obs = self._get_obs() if not self.done else None
return next_obs, reward_obj, self.done
def state(self):
return {
"current_task": self.current_task_index,
"current_step": self.current_step,
"partial_score": self.partial_score,
"done": self.done
} |