class UltrafeedbackEnvironment(Environment):
    def __init__(
        self,
        dataset_name: str = "trl-lib/ultrafeedback_binarized",
        split: str = "train",
        seed: int | None = None,
        reward_fn: callable = None,
    ):
        self.dataset_name = dataset_name
        self.split = split
        self.reward_fn = reward_fn
 
        # Load and format dataset
        self._load_dataset()
 
        # Initialize state
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._current_sample: dict | None = None

SYSTEM_PROMPT = "You are presented with a user prompt and a response from an LLM assistant. Grade the response to the prompt as a number between 0 and 10. DO NOT provide an answer to the original human prompt. The answer should be formatted as <answer>{number}</answer>."
 
def _load_dataset(self):
    """Load and format the dataset."""
    dataset = load_dataset(self.dataset_name, split=self.split)
    system_message = {"role": "system", "content": SYSTEM_PROMPT}
    new_samples = {"prompt": [], "answer": []}
 
    for sample in dataset:
        # Add chosen response
        new_samples["prompt"].append([system_message] + sample["chosen"])
        new_samples["answer"].append(sample["score_chosen"])
        # Add rejected response
        new_samples["prompt"].append([system_message] + sample["rejected"])
        new_samples["answer"].append(sample["score_rejected"])
 
    new_dataset = Dataset.from_dict(new_samples)
    # Cast answer to string
    new_dataset = new_dataset.cast_column("answer", Value("string"))

def step(self, action: UltrafeedbackAction) -> UltrafeedbackObservation:
    if self._current_sample is None:
        raise RuntimeError("Environment not initialized. Call reset() first.")
 
    self._state.step_count += 1
 
    completion = action.completion
    correct_answer_str = self._current_sample["answer"]
 
    response = parse_numeric_answer(completion)
 
    if response is None:
        correctness_reward = 0.0
    else:
        correct_answer = float(correct_answer_str)
        difference = abs(response - correct_answer)
        correctness_reward = 1.0 - (difference / 10.0)
        correctness_reward = max(0.0, correctness_reward)
 
    format_reward = calculate_format_reward(completion)
    total_reward = (correctness_reward * 1.0 + format_reward * 0.2) / 1.2
 
    return UltrafeedbackObservation(
        prompt=self._current_sample["prompt"],
        correct_answer=correct_answer_str,
        done=True,
        reward=total_reward,
        metadata={
            "predicted_score": response,
            "correct_score": correct_answer_str,
            "correctness_reward": correctness_reward,
            "format_reward": format_reward,
        },
    )

env = UltrafeedbackEnvironment()
 
# Hack: `GRPOTrainer` still expects a dataset argument, but we handle dataset logic in our env
train_dataset = Dataset.from_dict(
{
    "prompt": [
        "Blah blah this will be ignored"
    ]
    * 64
}
)

def reward_from_env(
    prompts: list[str], completions: list[str], **kwargs
) -> list[float]:
    """Extract environment rewards."""
    env_rewards = kwargs.get("raw_rewards", [])
    return (
        [float(reward) for reward in env_rewards]
        if env_rewards
        else [0.0] * len(completions)
    )
 
trainer = GRPOTrainer(
    model=model_name,
    reward_funcs=reward_from_env,
    train_dataset=train_dataset,
    rollout_func=rollout_func,
    processing_class=tokenizer,
    args=GRPOConfig(
        output_dir=str(Path(f"/models/{experiment_name}")),
        vllm_mode="server",
        use_vllm=True,
        num_train_epochs=1,
        num_generations=4,
        max_completion_length=128,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=1,
        logging_steps=1,
        report_to="wandb",
        temperature=0.8,
        top_k=10,
    ),
)
 
trainer.train()

app = modal.App(name="trl-openenv")
 
PYPROJECT_DEPENDENCIES = [
    "datasets>=4.4.1",
    "transformers>=4.57.1",
    "trl[vllm]>=0.25.1",
    "wandb>=0.22.3",
    "weave",
]
 
image = (
    modal.Image.debian_slim(python_version="3.12")
    .uv_pip_install(*PYPROJECT_DEPENDENCIES)
    # I have OpenEnv as a git subtree so Cursor can easily reference it, you could also install from Github
    .add_local_dir("OpenEnv", remote_path="/repos/OpenEnv", copy=True)
    .uv_pip_install("/repos/OpenEnv")
    # Add ultrafeedback environment
    .add_local_python_source("ultrafeedback", copy=True)
    .add_local_dir("src", remote_path="/root/src")
)
 
def _train_grpo(model_name: str, experiment_name: str):
    start_grpo_trainer(model_name, experiment_name)
 
 
@app.function(
    image=image,
    gpu="A10G:2",
    timeout=60 * 60 * 24,  # 24 hours
    secrets=[modal.Secret.from_name("wandb-secret")],
    volumes={str(MODELS_DIR): checkpoints_volume},
)
def train_grpo(model_name: str, experiment_name: str):
    experiment_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    _train_grpo(model_name, f"{experiment_name}_{experiment_timestamp}")
 

def start_grpo_trainer():
    env_copy = os.environ.copy()
    env_copy["CUDA_VISIBLE_DEVICES"] = "0"  # Run serve vLLM process on GPU 0
 
    # Start vllm-serve in the background
    vllm_process = subprocess.Popen(
        [
            "trl",
            "vllm-serve",
            "--model",
            model_name,
            "--gpu-memory-utilization",
            "0.8",
        ],
        env=env_copy,
    )
 
    # Make sure to run training on our other GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
    trainer = ...