Reward System Implementation

from atlas_core.training.algorithms.grpo import GRPOTrainer
from atlas_core.reward.interpretation import RIMReward
from datasets import load_dataset

# 1. Instantiate reward system
reward_system = RIMReward(config_path='reward_system/interpretation.yaml')

# 2. Pass to trainer
trainer = GRPOTrainer(
    model="path/to/your/teacher_model",
    args=grpo_config,
    reward_funcs=[reward_system],  # Just pass it in
    train_dataset=train_dataset
)

# 3. Train - the reward system runs automatically
trainer.train()

from atlas_core.reward.interpretation import RIMReward

# Create reward system
reward = RIMReward(config_path='reward_system/interpretation.yaml')

# Evaluate a single interaction
result = reward.evaluate(
    prompt="What is 2+2?",
    response="The answer is <solution>4</solution>.",
    baseline_solutions="It is 4",
    teacher_traces="Explain your reasoning step by step",
)

print(f"Score: {result.score}")
print(f"Per-judge: {result.judge_scores}")
print(f"Rationale:\\n{result.rationale}")

# src/atlas_core/reward/interpretation/judges.py
class AccuracyJudge:
    def _build_prompt(self, inputs: Dict[str, Any]) -> str:
        # Customize this string to change evaluation criteria
        return f"""Evaluate these responses.

Prompt: {inputs.get('prompt', '')}
Response A: {inputs.get('response_a', '')}
Response B: {inputs.get('response_b', '')}

Step 1: Generate 2-3 evaluation principles with weights (must sum to 1.0)
Step 2: Score both responses against each principle
Step 3: Provide final scores (0.0 to 1.0)

Output JSON only: {{"principles": [...], "score_a": float, "score_b": float, "uncertainty": float}}"""

class CreativityJudge:
    def __init__(self):
        self.name = 'creativity'

    def evaluate(self, inputs: Dict[str, Any], model_fn, temperature: float):
        prompt = f"""Score creativity (0.0 = formulaic, 1.0 = highly creative).
        Response: {inputs.get('response', '')}
        Output JSON: {{"score": float, "rationale": str, "uncertainty": float}}"""

        response = model_fn(prompt, temperature)
        return json.loads(response)

from atlas_core.reward.interpretation.judges import AccuracyJudge, HelpfulnessJudge, CreativityJudge

class RIMReward:
    def __init__(self, ...):
        self.judges = {
            'accuracy': AccuracyJudge(),
            'helpfulness': HelpfulnessJudge(),
            'creativity': CreativityJudge()  # Add here
        }

active_judges:
  accuracy: true
  helpfulness: true
  creativity: true  # Enable new judge

# Example log entry
{
  'step': 150,
  'rim_rewards': {
    'accuracy': 0.85,
    'helpfulness': 0.72,
    'process': 0.78,
    'diagnostic': 0.80
  },
  'rim_explanations': {
    'accuracy': 'Response correctly solves the problem with proper units',
    'helpfulness': 'Teaching improved reasoning structure significantly'
  },
  'escalation_rate': 0.23  # 23% of cases went to Tier 2
}