def reward_fn(completion, **kwargs):
    """
    Evaluate the model response and return a reward score.

    Args:
        completion: The model's response, formatted as list of messages
            e.g. [{"role": "assistant", "content": "This is an example reponse."}]
        kwargs: Additional arguments given by any keywords other than "prompt" in the prompts.jsonl file

    Returns:
        float: Reward score (higher is better)
    """
    # To do this the jsonl must have an "expected_result" field for every line
    expected = kwargs['expected_result']
    response = completion[0]["content"]

    # Example: Check if the response contains expected answer
    if expected and str(expected) in response:
        return 1.0

    return 0.5

def extract_answer(str):
    # Extracts answer from between <answer> tags
    if "<answer>" in str and "</answer>" in str:
        return str.split("<answer>")[1].split("</answer>")[0]
    else:
        return ""

def response_penalty(str):
    # Penalizes use of digits in the response
    return sum(1 for c in str if c.isdigit())

def reward_fn(completion, **kwargs):
    # Get expected result from kwargs
    expected_result = kwargs.get('expected_result')
    if expected_result is None:
        raise ValueError("expected_result is required for this reward function")

    # Get model response from completion
    model_response = completion[0]['content']

    # Extract answer from response
    extracted_answer = extract_answer(model_response)

    # Calculate correctness reward (10.0 for correct, 0.0 for incorrect)
    response_correctness_reward = 10.0 if extracted_answer == str(expected_result) else 0.0

    # Calculate penalty for using digits
    response_penalty = -response_penalty(model_response)

    # Combine reward and penalty
    reward = response_correctness_reward + response_penalty

    return reward

class CustomEnv:
    def setup(self, **kwargs):
        # Initialize your environment state here
        # kwargs contains any parameters from the JSONL file
        self.max_steps = kwargs.get("max_steps", 10)
        self.current_step = 0
        self.state = None

    def step(self, action: str):
        # Process the action and return the result
        self.current_step += 1

        # Your environment logic here
        # action is the model's output as a string

        done = self.current_step >= self.max_steps

        return {
            "observation": "Environment response to action",
            "reward": 0.0,  # Float reward value
            "done": done    # Boolean indicating episode end
        }

class MathProblemEnv:
    def setup(self, **kwargs):
        self.problem = kwargs.get("problem", "2 + 2")
        self.expected_answer = kwargs.get("expected_answer", "4")
        self.max_attempts = kwargs.get("max_attempts", 3)
        self.current_attempt = 0
        self.solved = False

    def step(self, action: str):
        self.current_attempt += 1

        # Check if the action contains the correct answer
        if str(self.expected_answer) in action:
            self.solved = True
            return {
                "observation": "Correct! You solved the problem.",
                "reward": 10.0,
                "done": True
            }

        # Check if max attempts reached
        if self.current_attempt >= self.max_attempts:
            return {
                "observation": f"Maximum attempts reached. The answer was {self.expected_answer}.",
                "reward": 0.0,
                "done": True
            }

        # Provide hint for next attempt
        remaining = self.max_attempts - self.current_attempt
        return {
            "observation": f"Incorrect. You have {remaining} attempt(s) remaining. Try again.",
            "reward": -1.0,
            "done": False
        }