Spaces:

testtej44
/

Code_man

Runtime error

Code_man / inference_backup.py

Tejasva

fix: inference_backup.py - strict START/STEP/END compliance

e38501f 7 months ago

16.5 kB

#!/usr/bin/env python3

import os
import json
import logging
from typing import Dict, Any, Optional, List

try:
    from openai import OpenAI
except ImportError:
    print("Error: openai package not found. Install with: pip install openai")
    exit(1)

from env.environment import CodeGuardianEnv
from env.models import Action, AgentResponse, AgentState

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def setup_openai_client() -> tuple[OpenAI, str]:
    """Initialize OpenAI client from environment variables."""
    api_base_url = os.getenv('API_BASE_URL', 'https://api.openai.com/v1')
    model_name = os.getenv('MODEL_NAME', 'gpt-4')
    
    # HF_TOKEN without default (required for Hugging Face)
    hf_token = os.getenv('HF_TOKEN')
    api_key = hf_token or os.getenv('API_KEY')
    
    if not api_key:
        raise ValueError("Neither HF_TOKEN nor API_KEY environment variables are set")
    
    # Initialize OpenAI client with base_url and api_key
    client = OpenAI(
        api_key=api_key,
        base_url=api_base_url
    )
    
    return client, model_name

def create_adaptive_prompt(observation: Dict[str, Any], step_type: str, agent_state: AgentState) -> str:
    """Create adaptive prompt based on agent learning state."""
    step_names = {
        'detect': 'Issue Detection',
        'risk': 'Risk Assessment', 
        'fix': 'Fix Suggestion'
    }
    
    step_instructions = {
        'detect': 'list all issues you find in the code',
        'risk': 'output the risk level (low/medium/high)',
        'fix': 'provide corrected code that fixes the identified issues'
    }
    
    # Adaptive prompt based on learning
    learning_context = ""
    if agent_state.past_failures:
        learning_context = f"\nLEARNING FROM PAST MISTAKES:\n{chr(10).join(f'- Avoid: {failure}' for failure in agent_state.past_failures[-3:])}"
    
    if agent_state.prompt_strategy == "detailed":
        instruction_detail = " Be thorough and provide detailed reasoning."
    elif agent_state.prompt_strategy == "focused":
        instruction_detail = " Focus on most critical issues first."
    else:
        instruction_detail = ""
    
    prompt = f"""You are a professional code reviewer learning from feedback.
Code:
{observation.code}
Intent:
{observation.intent}
Constraints:
{chr(10).join(f"- {constraint}" for constraint in observation.constraints)}
Step: {step_names[step_type]}
Instructions:
* {step_instructions[step_type]}{instruction_detail}
* Provide confidence score (0.0-1.0)
* Include your reasoning{learning_context}
Respond ONLY in JSON format:
{{
  "issues_detected": [...],
  "risk_level": "...",
  "fix": "...",
  "confidence": 0.0,
  "reasoning": "...",
  "reward_breakdown": {{
    "detect": 0.0,
    "risk": 0.0,
    "fix": 0.0
  }}
}}"""
    
    return prompt

def parse_llm_response(response_text: str) -> AgentResponse:
    """Parse LLM response into structured AgentResponse with fallback."""
    try:
        # Clean up response text
        response_text = response_text.strip()
        if response_text.startswith('```json'):
            response_text = response_text[7:]
        if response_text.endswith('```'):
            response_text = response_text[:-3]
        response_text = response_text.strip()
        
        parsed = json.loads(response_text)
        
        # Ensure all required fields exist with defaults
        return AgentResponse(
            issues_detected=parsed.get('issues_detected', []),
            risk_level=parsed.get('risk_level', 'low'),
            fix=parsed.get('fix', ''),
            confidence=max(0.0, min(1.0, parsed.get('confidence', 0.5))),
            reasoning=parsed.get('reasoning', ''),
            reward_breakdown=parsed.get('reward_breakdown', {})
        )
    except (json.JSONDecodeError, AttributeError, ValueError) as e:
        logger.warning(f"Failed to parse LLM response: {e}")
        logger.warning(f"Raw response: {response_text[:200]}...")
        
        # Return fallback response
        return AgentResponse(
            issues_detected=[],
            risk_level='low',
            fix='',
            confidence=0.1,
            reasoning='Failed to parse response',
            reward_breakdown={}
        )

def create_action_from_response(step_type: str, response: AgentResponse) -> Action:
    """Create Action model from structured AgentResponse."""
    return Action(
        action_type=step_type,
        detected_issues=response.issues_detected if step_type == 'detect' else None,
        risk_level=response.risk_level if step_type == 'risk' else None,
        suggested_fix=response.fix if step_type == 'fix' else None
    )

def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
    """Run LLM inference with given prompt."""
    try:
        response = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": "You are a professional code reviewer. Always respond in valid JSON format with confidence and reasoning."},
                {"role": "user", "content": prompt}
            ],
            temperature=0,  # Deterministic
            max_tokens=500  # Increased for structured output
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"LLM inference failed: {e}")
        raise

def run_baseline(env: CodeGuardianEnv) -> float:
    """Run a simple non-learning baseline agent for comparison."""
    print("STEP: Running baseline agent")
    logger.info("[BASELINE START] Running baseline agent...")
    
    obs = env.reset()
    total_reward = 0.0
    step_count = 0
    
    while not env.done and step_count < 50:
        current_step = env.state()['current_step']
        step_types = ['detect', 'risk', 'fix']
        step_type = step_types[current_step]
        
        # Baseline always outputs minimal responses
        if step_type == 'detect':
            action = Action(action_type='detect', detected_issues=[])
        elif step_type == 'risk':
            action = Action(action_type='risk', risk_level='low')
        else:
            action = Action(action_type='fix', suggested_fix='')
        
        next_obs, reward, done = env.step(action)
        total_reward += reward.value
        step_count += 1
        obs = next_obs
        
        logger.info(f"[BASELINE STEP] {step_count}: {step_type} -> reward={reward.value:+.3f}")
    
    print(f"STEP: Baseline complete - total_reward={total_reward:.3f} steps={step_count}")
    logger.info(f"[BASELINE END] Total reward: {total_reward:.3f}, Steps: {step_count}")
    return total_reward

def update_agent_state(agent_state: AgentState, episode_reward: float, feedback: str) -> AgentState:
    """Update agent state based on episode performance for policy improvement."""
    agent_state.episode += 1
    agent_state.total_reward = episode_reward
    agent_state.improvement_history.append(episode_reward)
    
    # Update best score
    if episode_reward > agent_state.best_score:
        agent_state.best_score = episode_reward
        logger.info(f"New best score achieved: {episode_reward:.3f}")
    
    # STRONGER POLICY UPDATE LOGIC
    # Treat poor performance as failure
    if episode_reward < agent_state.best_score * 0.8:
        agent_state.past_failures.append(f"Episode {agent_state.episode}: reward={episode_reward:.3f}")
        # Keep only recent failures
        agent_state.past_failures = agent_state.past_failures[-5:]
        logger.info(f"Failure recorded - Episode reward: {episode_reward:.3f} vs Best: {agent_state.best_score:.3f}")
    
    # IMPROVED STRATEGY SWITCHING
    if episode_reward < agent_state.best_score * 0.7:  # Poor performance threshold
        old_strategy = agent_state.prompt_strategy
        
        if agent_state.prompt_strategy == "default":
            agent_state.prompt_strategy = "detailed"
            logger.info("Switching strategy: default → detailed")
        elif agent_state.prompt_strategy == "detailed":
            agent_state.prompt_strategy = "focused"
            logger.info("Switching strategy: detailed → focused")
        else:
            # Reset to default if focused still fails
            agent_state.prompt_strategy = "default"
            logger.info("Switching strategy: focused → default (reset)")
        
        if old_strategy != agent_state.prompt_strategy:
            logger.info(f"Strategy change triggered by low performance: {episode_reward:.3f} < {agent_state.best_score * 0.7:.3f}")
    
    return agent_state

def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_state: AgentState) -> tuple[float, List[Dict]]:
    """Run a single episode and return total reward with episode history."""
    obs = env.reset()
    episode_reward = 0.0
    episode_history = []
    step_count = 0
    
    print(f"START: Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
    logger.info(f"[EPISODE START] Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
    
    while not env.done:
        current_step = env.state()['current_step']
        step_types = ['detect', 'risk', 'fix']
        step_type = step_types[current_step]
        
        step_count += 1
        print(f"STEP: {step_count} - {step_type} - task_id={obs.task_id}")
        logger.info(f"[STEP] step_number={step_count} step_type={step_type} task_id={obs.task_id}")
        
        # Create adaptive prompt
        prompt = create_adaptive_prompt(obs, step_type, agent_state)
        
        # Run LLM inference
        try:
            response_text = run_llm_inference(client, model_name, prompt)
            logger.info(f"LLM response: {response_text[:100]}...")
        except Exception as e:
            logger.error(f"LLM inference failed: {e}")
            # Use fallback response
            agent_response = AgentResponse(
                issues_detected=[],
                risk_level='low',
                fix='',
                confidence=0.1,
                reasoning='Inference failed',
                reward_breakdown={}
            )
        else:
            # Parse structured response
            agent_response = parse_llm_response(response_text)
        
        # Create action from response
        action = create_action_from_response(step_type, agent_response)
        
        # Execute step
        next_obs, reward, done = env.step(action)
        
        # CONFIDENCE-BASED REWARD ADJUSTMENT
        confidence_penalty = (1 - agent_response.confidence) * 0.1
        adjusted_reward = reward.value - confidence_penalty
        episode_reward += adjusted_reward
        
        logger.info(f"confidence_penalty={confidence_penalty:.3f} adjusted_reward={adjusted_reward:+.3f} confidence={agent_response.confidence:.2f}")
        
        # Log step results
        step_info = {
            "step": step_count,
            "type": step_type,
            "action": action.dict(),
            "response": agent_response.dict(),
            "reward": reward.dict(),
            "adjusted_reward": adjusted_reward,
            "confidence_penalty": confidence_penalty,
            "total_reward": episode_reward
        }
        episode_history.append(step_info)
        
        logger.info(f"action_type={action.action_type} reward={reward.value:+.3f} adjusted_reward={adjusted_reward:+.3f} feedback={reward.feedback} total_reward={episode_reward:+.3f} current_step={env.state()['current_step']} current_task={env.state()['current_task']}")
        
        # Update observation for next step
        obs = next_obs
        
        # Safety check
        if step_count > 50:  # Prevent infinite loops
            logger.warning("Maximum steps reached, stopping episode")
            break
    
    # Final episode log
    print(f"END: Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
    logger.info(f"[EPISODE END] Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
    
    return episode_reward, episode_history

def train_agent(num_episodes: int = 10) -> AgentState:
    """Train agent over multiple episodes with policy improvement."""
    try:
        # Setup
        client, model_name = setup_openai_client()
        agent_state = AgentState()
        all_episodes = []
        
        print(f"START: RL training for {num_episodes} episodes with model: {model_name}")
        logger.info(f"[TRAINING START] Starting RL training for {num_episodes} episodes with model: {model_name}")
        
        # RUN BASELINE FIRST
        logger.info("\n{'='*60}")
        logger.info("BASELINE EVALUATION")
        logger.info(f"{'='*60}")
        baseline_env = CodeGuardianEnv()
        baseline_score = run_baseline(baseline_env)
        logger.info(f"Baseline score: {baseline_score:.3f}")
        
        # TRAINING LOOP
        logger.info(f"\n{'='*60}")
        logger.info("REINFORCEMENT LEARNING")
        logger.info(f"{'='*60}")
        
        for episode in range(num_episodes):
            logger.info(f"\n{'='*50}")
            logger.info(f"EPISODE {episode + 1}/{num_episodes}")
            logger.info(f"{'='*50}")
            
            env = CodeGuardianEnv()
            episode_reward, episode_history = run_episode(env, client, model_name, agent_state)
            all_episodes.append({
                "episode": episode + 1,
                "reward": episode_reward,
                "history": episode_history
            })
            
            # Update agent state for policy improvement
            agent_state = update_agent_state(agent_state, episode_reward, f"Episode {episode + 1} completed")
            
            logger.info(f"Episode {episode + 1} completed with reward: {episode_reward:.3f}")
            logger.info(f"Best score so far: {agent_state.best_score:.3f}")
            logger.info(f"Current strategy: {agent_state.prompt_strategy}")
            
            # EARLY STOPPING
            if episode > 3 and agent_state.best_score > 2.5:
                logger.info("Early stopping triggered: high performance achieved")
                logger.info(f"Stopping at episode {episode + 1} with best score {agent_state.best_score:.3f}")
                break
        
        # PERFORMANCE COMPARISON
        print("STEP: Training summary analysis")
        logger.info(f"\n{'='*60}")
        logger.info(f"[TRAINING SUMMARY]")
        logger.info(f"{'='*60}")
        
        # Calculate improvement over baseline
        improvement = ((agent_state.best_score - baseline_score) / (abs(baseline_score) + 1e-6)) * 100
        
        # Log key metrics
        logger.info(f"RL Agent Best Score: {agent_state.best_score:.3f}")
        logger.info(f"Baseline Score: {baseline_score:.3f}")
        logger.info(f"Performance Improvement: {improvement:.2f}%")
        
        # Additional metrics
        if all_episodes:
            avg_reward = sum(ep["reward"] for ep in all_episodes) / len(all_episodes)
            logger.info(f"Average Episode Reward: {avg_reward:.3f}")
            logger.info(f"Episodes Completed: {len(all_episodes)}")
            logger.info(f"Final Strategy: {agent_state.prompt_strategy}")
            
            # Performance classification
            if improvement > 50:
                print("STEP: EXCELLENT - significant improvement over baseline")
            elif improvement > 20:
                print("STEP: GOOD - measurable improvement achieved")
            elif improvement > 0:
                print("STEP: POSITIVE - some improvement over baseline")
            else:
                print("STEP: NO IMPROVEMENT - RL agent not better than baseline")
        
        return agent_state
        
    except Exception as e:
        logger.error(f"Training failed: {e}")
        raise

def main():
    """Main inference loop with RL training."""
    try:
        print("START")
        
        # Run RL training
        agent_state = train_agent(num_episodes=10)
        return 0
        
    except Exception as e:
        logger.error(f"Inference failed: {e}")
        return 1
    finally:
        print("END")

if __name__ == "__main__":
    exit(main())

	#!/usr/bin/env python3

	import os
	import json
	import logging
	from typing import Dict, Any, Optional, List

	try:
	from openai import OpenAI
	except ImportError:
	print("Error: openai package not found. Install with: pip install openai")
	exit(1)

	from env.environment import CodeGuardianEnv
	from env.models import Action, AgentResponse, AgentState

	# Setup logging
	logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
	logger = logging.getLogger(__name__)

	def setup_openai_client() -> tuple[OpenAI, str]:
	"""Initialize OpenAI client from environment variables."""
	api_base_url = os.getenv('API_BASE_URL', 'https://api.openai.com/v1')
	model_name = os.getenv('MODEL_NAME', 'gpt-4')

	# HF_TOKEN without default (required for Hugging Face)
	hf_token = os.getenv('HF_TOKEN')
	api_key = hf_token or os.getenv('API_KEY')

	if not api_key:
	raise ValueError("Neither HF_TOKEN nor API_KEY environment variables are set")

	# Initialize OpenAI client with base_url and api_key
	client = OpenAI(
	api_key=api_key,
	base_url=api_base_url
	)

	return client, model_name

	def create_adaptive_prompt(observation: Dict[str, Any], step_type: str, agent_state: AgentState) -> str:
	"""Create adaptive prompt based on agent learning state."""
	step_names = {
	'detect': 'Issue Detection',
	'risk': 'Risk Assessment',
	'fix': 'Fix Suggestion'
	}

	step_instructions = {
	'detect': 'list all issues you find in the code',
	'risk': 'output the risk level (low/medium/high)',
	'fix': 'provide corrected code that fixes the identified issues'
	}

	# Adaptive prompt based on learning
	learning_context = ""
	if agent_state.past_failures:
	learning_context = f"\nLEARNING FROM PAST MISTAKES:\n{chr(10).join(f'- Avoid: {failure}' for failure in agent_state.past_failures[-3:])}"

	if agent_state.prompt_strategy == "detailed":
	instruction_detail = " Be thorough and provide detailed reasoning."
	elif agent_state.prompt_strategy == "focused":
	instruction_detail = " Focus on most critical issues first."
	else:
	instruction_detail = ""

	prompt = f"""You are a professional code reviewer learning from feedback.

	Code:
	{observation.code}

	Intent:
	{observation.intent}

	Constraints:
	{chr(10).join(f"- {constraint}" for constraint in observation.constraints)}

	Step: {step_names[step_type]}

	Instructions:
	* {step_instructions[step_type]}{instruction_detail}
	* Provide confidence score (0.0-1.0)
	* Include your reasoning{learning_context}

	Respond ONLY in JSON format:
	{{
	"issues_detected": [...],
	"risk_level": "...",
	"fix": "...",
	"confidence": 0.0,
	"reasoning": "...",
	"reward_breakdown": {{
	"detect": 0.0,
	"risk": 0.0,
	"fix": 0.0
	}}
	}}"""

	return prompt

	def parse_llm_response(response_text: str) -> AgentResponse:
	"""Parse LLM response into structured AgentResponse with fallback."""
	try:
	# Clean up response text
	response_text = response_text.strip()
	if response_text.startswith('```json'):
	response_text = response_text[7:]
	if response_text.endswith('```'):
	response_text = response_text[:-3]
	response_text = response_text.strip()

	parsed = json.loads(response_text)

	# Ensure all required fields exist with defaults
	return AgentResponse(
	issues_detected=parsed.get('issues_detected', []),
	risk_level=parsed.get('risk_level', 'low'),
	fix=parsed.get('fix', ''),
	confidence=max(0.0, min(1.0, parsed.get('confidence', 0.5))),
	reasoning=parsed.get('reasoning', ''),
	reward_breakdown=parsed.get('reward_breakdown', {})
	)
	except (json.JSONDecodeError, AttributeError, ValueError) as e:
	logger.warning(f"Failed to parse LLM response: {e}")
	logger.warning(f"Raw response: {response_text[:200]}...")

	# Return fallback response
	return AgentResponse(
	issues_detected=[],
	risk_level='low',
	fix='',
	confidence=0.1,
	reasoning='Failed to parse response',
	reward_breakdown={}
	)

	def create_action_from_response(step_type: str, response: AgentResponse) -> Action:
	"""Create Action model from structured AgentResponse."""
	return Action(
	action_type=step_type,
	detected_issues=response.issues_detected if step_type == 'detect' else None,
	risk_level=response.risk_level if step_type == 'risk' else None,
	suggested_fix=response.fix if step_type == 'fix' else None
	)

	def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
	"""Run LLM inference with given prompt."""
	try:
	response = client.chat.completions.create(
	model=model_name,
	messages=[
	{"role": "system", "content": "You are a professional code reviewer. Always respond in valid JSON format with confidence and reasoning."},
	{"role": "user", "content": prompt}
	],
	temperature=0, # Deterministic
	max_tokens=500 # Increased for structured output
	)
	return response.choices[0].message.content
	except Exception as e:
	logger.error(f"LLM inference failed: {e}")
	raise

	def run_baseline(env: CodeGuardianEnv) -> float:
	"""Run a simple non-learning baseline agent for comparison."""
	print("STEP: Running baseline agent")
	logger.info("[BASELINE START] Running baseline agent...")

	obs = env.reset()
	total_reward = 0.0
	step_count = 0

	while not env.done and step_count < 50:
	current_step = env.state()['current_step']
	step_types = ['detect', 'risk', 'fix']
	step_type = step_types[current_step]

	# Baseline always outputs minimal responses
	if step_type == 'detect':
	action = Action(action_type='detect', detected_issues=[])
	elif step_type == 'risk':
	action = Action(action_type='risk', risk_level='low')
	else:
	action = Action(action_type='fix', suggested_fix='')

	next_obs, reward, done = env.step(action)
	total_reward += reward.value
	step_count += 1
	obs = next_obs

	logger.info(f"[BASELINE STEP] {step_count}: {step_type} -> reward={reward.value:+.3f}")

	print(f"STEP: Baseline complete - total_reward={total_reward:.3f} steps={step_count}")
	logger.info(f"[BASELINE END] Total reward: {total_reward:.3f}, Steps: {step_count}")
	return total_reward

	def update_agent_state(agent_state: AgentState, episode_reward: float, feedback: str) -> AgentState:
	"""Update agent state based on episode performance for policy improvement."""
	agent_state.episode += 1
	agent_state.total_reward = episode_reward
	agent_state.improvement_history.append(episode_reward)

	# Update best score
	if episode_reward > agent_state.best_score:
	agent_state.best_score = episode_reward
	logger.info(f"New best score achieved: {episode_reward:.3f}")

	# STRONGER POLICY UPDATE LOGIC
	# Treat poor performance as failure
	if episode_reward < agent_state.best_score * 0.8:
	agent_state.past_failures.append(f"Episode {agent_state.episode}: reward={episode_reward:.3f}")
	# Keep only recent failures
	agent_state.past_failures = agent_state.past_failures[-5:]
	logger.info(f"Failure recorded - Episode reward: {episode_reward:.3f} vs Best: {agent_state.best_score:.3f}")

	# IMPROVED STRATEGY SWITCHING
	if episode_reward < agent_state.best_score * 0.7: # Poor performance threshold
	old_strategy = agent_state.prompt_strategy

	if agent_state.prompt_strategy == "default":
	agent_state.prompt_strategy = "detailed"
	logger.info("Switching strategy: default → detailed")
	elif agent_state.prompt_strategy == "detailed":
	agent_state.prompt_strategy = "focused"
	logger.info("Switching strategy: detailed → focused")
	else:
	# Reset to default if focused still fails
	agent_state.prompt_strategy = "default"
	logger.info("Switching strategy: focused → default (reset)")

	if old_strategy != agent_state.prompt_strategy:
	logger.info(f"Strategy change triggered by low performance: {episode_reward:.3f} < {agent_state.best_score * 0.7:.3f}")

	return agent_state

	def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_state: AgentState) -> tuple[float, List[Dict]]:
	"""Run a single episode and return total reward with episode history."""
	obs = env.reset()
	episode_reward = 0.0
	episode_history = []
	step_count = 0

	print(f"START: Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
	logger.info(f"[EPISODE START] Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")

	while not env.done:
	current_step = env.state()['current_step']
	step_types = ['detect', 'risk', 'fix']
	step_type = step_types[current_step]

	step_count += 1
	print(f"STEP: {step_count} - {step_type} - task_id={obs.task_id}")
	logger.info(f"[STEP] step_number={step_count} step_type={step_type} task_id={obs.task_id}")

	# Create adaptive prompt
	prompt = create_adaptive_prompt(obs, step_type, agent_state)

	# Run LLM inference
	try:
	response_text = run_llm_inference(client, model_name, prompt)
	logger.info(f"LLM response: {response_text[:100]}...")
	except Exception as e:
	logger.error(f"LLM inference failed: {e}")
	# Use fallback response
	agent_response = AgentResponse(
	issues_detected=[],
	risk_level='low',
	fix='',
	confidence=0.1,
	reasoning='Inference failed',
	reward_breakdown={}
	)
	else:
	# Parse structured response
	agent_response = parse_llm_response(response_text)

	# Create action from response
	action = create_action_from_response(step_type, agent_response)

	# Execute step
	next_obs, reward, done = env.step(action)

	# CONFIDENCE-BASED REWARD ADJUSTMENT
	confidence_penalty = (1 - agent_response.confidence) * 0.1
	adjusted_reward = reward.value - confidence_penalty
	episode_reward += adjusted_reward

	logger.info(f"confidence_penalty={confidence_penalty:.3f} adjusted_reward={adjusted_reward:+.3f} confidence={agent_response.confidence:.2f}")

	# Log step results
	step_info = {
	"step": step_count,
	"type": step_type,
	"action": action.dict(),
	"response": agent_response.dict(),
	"reward": reward.dict(),
	"adjusted_reward": adjusted_reward,
	"confidence_penalty": confidence_penalty,
	"total_reward": episode_reward
	}
	episode_history.append(step_info)

	logger.info(f"action_type={action.action_type} reward={reward.value:+.3f} adjusted_reward={adjusted_reward:+.3f} feedback={reward.feedback} total_reward={episode_reward:+.3f} current_step={env.state()['current_step']} current_task={env.state()['current_task']}")

	# Update observation for next step
	obs = next_obs

	# Safety check
	if step_count > 50: # Prevent infinite loops
	logger.warning("Maximum steps reached, stopping episode")
	break

	# Final episode log
	print(f"END: Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
	logger.info(f"[EPISODE END] Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")

	return episode_reward, episode_history

	def train_agent(num_episodes: int = 10) -> AgentState:
	"""Train agent over multiple episodes with policy improvement."""
	try:
	# Setup
	client, model_name = setup_openai_client()
	agent_state = AgentState()
	all_episodes = []

	print(f"START: RL training for {num_episodes} episodes with model: {model_name}")
	logger.info(f"[TRAINING START] Starting RL training for {num_episodes} episodes with model: {model_name}")

	# RUN BASELINE FIRST
	logger.info("\n{'='*60}")
	logger.info("BASELINE EVALUATION")
	logger.info(f"{'='*60}")
	baseline_env = CodeGuardianEnv()
	baseline_score = run_baseline(baseline_env)
	logger.info(f"Baseline score: {baseline_score:.3f}")

	# TRAINING LOOP
	logger.info(f"\n{'='*60}")
	logger.info("REINFORCEMENT LEARNING")
	logger.info(f"{'='*60}")

	for episode in range(num_episodes):
	logger.info(f"\n{'='*50}")
	logger.info(f"EPISODE {episode + 1}/{num_episodes}")
	logger.info(f"{'='*50}")

	env = CodeGuardianEnv()
	episode_reward, episode_history = run_episode(env, client, model_name, agent_state)
	all_episodes.append({
	"episode": episode + 1,
	"reward": episode_reward,
	"history": episode_history
	})

	# Update agent state for policy improvement
	agent_state = update_agent_state(agent_state, episode_reward, f"Episode {episode + 1} completed")

	logger.info(f"Episode {episode + 1} completed with reward: {episode_reward:.3f}")
	logger.info(f"Best score so far: {agent_state.best_score:.3f}")
	logger.info(f"Current strategy: {agent_state.prompt_strategy}")

	# EARLY STOPPING
	if episode > 3 and agent_state.best_score > 2.5:
	logger.info("Early stopping triggered: high performance achieved")
	logger.info(f"Stopping at episode {episode + 1} with best score {agent_state.best_score:.3f}")
	break

	# PERFORMANCE COMPARISON
	print("STEP: Training summary analysis")
	logger.info(f"\n{'='*60}")
	logger.info(f"[TRAINING SUMMARY]")
	logger.info(f"{'='*60}")

	# Calculate improvement over baseline
	improvement = ((agent_state.best_score - baseline_score) / (abs(baseline_score) + 1e-6)) * 100

	# Log key metrics
	logger.info(f"RL Agent Best Score: {agent_state.best_score:.3f}")
	logger.info(f"Baseline Score: {baseline_score:.3f}")
	logger.info(f"Performance Improvement: {improvement:.2f}%")

	# Additional metrics
	if all_episodes:
	avg_reward = sum(ep["reward"] for ep in all_episodes) / len(all_episodes)
	logger.info(f"Average Episode Reward: {avg_reward:.3f}")
	logger.info(f"Episodes Completed: {len(all_episodes)}")
	logger.info(f"Final Strategy: {agent_state.prompt_strategy}")

	# Performance classification
	if improvement > 50:
	print("STEP: EXCELLENT - significant improvement over baseline")
	elif improvement > 20:
	print("STEP: GOOD - measurable improvement achieved")
	elif improvement > 0:
	print("STEP: POSITIVE - some improvement over baseline")
	else:
	print("STEP: NO IMPROVEMENT - RL agent not better than baseline")

	return agent_state

	except Exception as e:
	logger.error(f"Training failed: {e}")
	raise

	def main():
	"""Main inference loop with RL training."""
	try:
	print("START")

	# Run RL training
	agent_state = train_agent(num_episodes=10)
	return 0

	except Exception as e:
	logger.error(f"Inference failed: {e}")
	return 1
	finally:
	print("END")

	if __name__ == "__main__":
	exit(main())