Spaces:

testtej44
/

Code_man

Runtime error

Code_man / env /environment.py

Tejasva

Deploy Codeman RL system

961dad9 7 months ago

3.72 kB

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
        }

	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
	}