From d761badb5bb6e7426e73c91d57441f4da011a96e Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Mon, 20 Apr 2026 17:28:02 +0200 Subject: [PATCH 1/6] docs(tutorials): add end-to-end Rubrics tutorial MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rubrics (RFC 004, `openenv.core.rubrics`) shipped with the OpenEnv core package but had no tutorial and were never introduced to readers outside the RFC. As a result the page at `guides/rewards.md` still described ad-hoc `compute_reward(obs, action, terminated)` functions — a pattern that contradicts the "rewards live inside the environment" invariant and the composable Rubric API. This PR: - Adds `docs/source/tutorials/rubrics.md`, an end-to-end walkthrough of the Rubric API. Covers the base class and hooks; the four containers (`Sequential`, `Gate`, `WeightedSum`, `RubricList` / `RubricDict`); `LLMJudge` with both hosted (OpenAI / Anthropic via `create_llm_client`) and open-weight paths (`OpenAIClient` against a vLLM / Ollama / Hugging Face Inference Providers endpoint); `TrajectoryRubric` and `ExponentialDiscountingTrajectoryRubric` for delayed rewards; and a minimal `CodeEnvironment` that shows the full wiring through `super().__init__(rubric=...)`, `self._reset_rubric()` in `reset`, and `self._apply_rubric(action, observation)` in `step`. A final subsection explains where the reward ends up during training — TRL's recommended `environment_factory` path, plus torchforge — and how `named_rubrics()` is used orthogonally for per-component logging (W&B / TensorBoard / trackio). - Rewrites `docs/source/guides/rewards.md` around rubrics. The three design principles (start simple, shape carefully, consider density) are preserved, but the examples are now expressed as rubrics and the page links out to the tutorial for the implementation details. - Introduces rubrics in the core docs so the tutorial has a natural entry point from the nav: * `docs/source/concepts.md` gets a short `### Rubric` section in "Key Abstractions", between `StepResult` and `Client`. * `docs/source/guides/environment-anatomy.md` gets a new "Rewards via the Rubric" section covering the `rubric` parameter on `Environment.__init__` and the `_apply_rubric` / `_reset_rubric` helpers. - Hooks the tutorial into `docs/source/tutorials/index.md`. The follow-up work to migrate the Wordle GRPO tutorial from `rollout_func` to `environment_factory` (the TRL-recommended path) is tracked separately; this PR links readers directly to TRL's OpenEnv integration guide for the current training recipe. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/source/concepts.md | 4 + docs/source/guides/environment-anatomy.md | 10 + docs/source/guides/rewards.md | 105 ++++--- docs/source/tutorials/index.md | 2 + docs/source/tutorials/rubrics.md | 354 ++++++++++++++++++++++ 5 files changed, 425 insertions(+), 50 deletions(-) create mode 100644 docs/source/tutorials/rubrics.md diff --git a/docs/source/concepts.md b/docs/source/concepts.md index c7810d26ed..3a803bc542 100644 --- a/docs/source/concepts.md +++ b/docs/source/concepts.md @@ -54,6 +54,10 @@ A **StepResult** bundles together everything returned from a step: - `truncated`: Whether the episode was cut short - `info`: Additional metadata +### Rubric + +A **Rubric** is a composable unit of reward computation that lives inside the environment. Rubrics can be combined (`WeightedSum`, `Gate`, `Sequential`) to express multi-criteria rewards, support LLM judges for subjective criteria, and support delayed trajectory-based scoring. Every environment declares `self.rubric` in its constructor; `step` calls it via `self._apply_rubric(action, observation)` and the result flows back through `observation.reward`. See the [Rubrics tutorial](tutorials/rubrics.md) for the full API. + ### Client A **Client** is how you connect to and interact with an environment. OpenEnv provides both async and sync clients. diff --git a/docs/source/guides/environment-anatomy.md b/docs/source/guides/environment-anatomy.md index 1b9d2d21c4..647ed88cac 100644 --- a/docs/source/guides/environment-anatomy.md +++ b/docs/source/guides/environment-anatomy.md @@ -90,6 +90,16 @@ from openenv.server import create_server app = create_server(MyEnvironment) ``` +## Rewards via the Rubric + +Rewards in OpenEnv are computed **inside the environment** — they never come from outside the `step` call. The base `Environment` accepts an optional `rubric` parameter on `__init__` and stores it as `self.rubric`. The recommended flow is: + +1. Build the rubric and pass it to `super().__init__(rubric=...)` so `self.rubric` is set. +2. Call `self._reset_rubric()` at the start of `reset` so trajectory state does not leak between episodes. +3. Call `self._apply_rubric(action, observation)` inside `step` and attach the result to `observation.reward`. Use `self._apply_rubric_async(...)` from `step_async`. + +The rubric system (`openenv.core.rubrics`) ships with composable containers (`WeightedSum`, `Gate`, `Sequential`), an `LLMJudge` for subjective criteria, and `TrajectoryRubric` for delayed rewards. The [Rubrics tutorial](../tutorials/rubrics.md) covers the full API. + ## Next Steps - [Deployment](deployment.md) - Deploy your environment diff --git a/docs/source/guides/rewards.md b/docs/source/guides/rewards.md index 13a1a6a20e..5f217d498b 100644 --- a/docs/source/guides/rewards.md +++ b/docs/source/guides/rewards.md @@ -1,88 +1,93 @@ # Reward Design -:::{note} -Coming Soon - This page is under construction. +Good reward signals are the single biggest lever on RL training outcomes — and the hardest thing to get right. This page covers the design principles that apply in OpenEnv and points to the [Rubrics tutorial](../tutorials/rubrics.md) for the concrete implementation API. -Learn how to design +## Rewards Live Inside the Environment -Learn how to design effective reward functions for your OpenEnv environments. +The OpenEnv contract is that reward computation stays on the server side, inside `Environment.step`. The `openenv.core.rubrics` module exists to make that computation composable — think `torch.nn.Module` but for reward functions. An environment declares `self.rubric` in its constructor and the base class helper `self._apply_rubric(action, observation)` runs it every step. -## Overview +See the [Rubrics tutorial](../tutorials/rubrics.md) for the end-to-end walkthrough, including `WeightedSum`, `Gate`, `Sequential`, `LLMJudge`, and the trajectory rubrics used for delayed rewards like chess outcomes. -Reward functions are critical for RL training. They signal to your agent what behaviors are desirable. +## Design Principles -## Reward Design Principles +The principles below are independent of the rubric mechanics — they apply whether the reward is a single scalar or a composition of a dozen components. -### 1. Start Simple +### 1. Start simple -Begin with sparse rewards (success/failure) before adding shaped rewards: +Begin with a sparse success/failure signal. Only shape the reward after you've confirmed the agent can reach the positive signal at all. ```python -def compute_reward(observation, action, terminated): - if terminated and observation.success: - return 1.0 - elif terminated: - return -1.0 - return 0.0 +class WinLossRubric(Rubric): + def forward(self, action, observation) -> float: + if not observation.done: + return 0.0 + return 1.0 if observation.success else -1.0 ``` -### 2. Shape Carefully +### 2. Shape carefully -Add intermediate rewards to help learning, but avoid reward hacking: +Dense intermediate rewards speed up learning but invite reward hacking. Prefer adding shaping as a small-weighted *component* of a `WeightedSum`, so the dominant signal is still the real outcome. ```python -def compute_reward(observation, action, terminated): - reward = 0.0 - - # Progress reward - reward += 0.1 * observation.progress_delta - - # Success bonus - if terminated and observation.success: - reward += 10.0 - - return reward +reward = WeightedSum( + [WinLossRubric(), ProgressRubric()], + weights=[0.8, 0.2], +) ``` -### 3. Consider Density +### 3. Consider density + +- **Sparse** rewards (signal only on terminal transitions) are cleaner but slower to learn from. +- **Dense** rewards (signal every step) are faster but can push the agent into local optima that optimise the proxy instead of the goal. -Dense rewards (every step) speed learning but can cause local optima. Sparse rewards are cleaner but slower. +For long-horizon tasks with a delayed outcome, reach for `TrajectoryRubric` and its built-in per-step credit assignment instead of hand-crafting a dense proxy. -## Environment Examples +## Worked Examples -### Chess (Sparse) +### Chess — sparse, trajectory-based ```python -# Win: +1, Loss: -1, Draw: 0 -reward = result.observation.game_result +class ChessOutcomeRubric(ExponentialDiscountingTrajectoryRubric): + def score_trajectory(self, trajectory) -> float: + _, final_obs = trajectory[-1] + return final_obs.reward # +1 / 0 / -1 from the engine ``` -### Coding (Dense) +Rewards accumulate silently during the game, then the final step produces the outcome and `compute_step_rewards()` distributes it back with exponential discounting. This is exactly what `envs/chess_env/` ships with. + +### Coding — gated composition ```python -reward = 0.0 -if observation.tests_passed > prev_tests_passed: - reward += 0.5 * (observation.tests_passed - prev_tests_passed) -if observation.all_tests_passed: - reward += 5.0 +reward = Sequential( + Gate(CompilesRubric(), threshold=1.0), # must compile + Gate(TestsPassRubric(), threshold=0.5), # must pass at least half the tests + WeightedSum([TestsPassRubric(), StyleRubric()], [0.7, 0.3]), +) ``` -### TextArena (Mixed) +`Sequential` short-circuits to `0.0` the moment any child fails its gate, so expensive style / LLM judge calls never run on broken submissions. + +### Text tasks — mixed signal ```python -# Per-turn progress + final outcome -reward = observation.score_delta + (10.0 if observation.won else 0.0) +reward = WeightedSum( + [WinLossRubric(), PerTurnProgressRubric()], + weights=[0.7, 0.3], +) ``` +Per-turn progress helps the agent explore, but the terminal outcome still dominates the final score. + ## Common Pitfalls -1. **Reward hacking** - Agent finds unintended shortcuts -2. **Sparse rewards** - Agent never finds positive signal -3. **Conflicting signals** - Mixed incentives confuse learning +1. **Reward hacking.** The agent optimises the proxy instead of the goal. Prefer shaping as a small-weighted component rather than a large intermediate bonus. +2. **Sparse rewards that never fire.** If the agent cannot reach the success signal in practice, training stalls. Measure success rate on a random policy before relying on a sparse reward. +3. **Conflicting signals.** Two criteria pulling in opposite directions produce a flat optimisation landscape. If you catch yourself subtracting rewards to "cancel out" a bad behaviour, consider a `Gate` instead — make the bad case a hard zero. +4. **Component score drift.** Without component-level logging, you won't know which criterion dropped the total. Rubric introspection (`env.rubric.named_rubrics()` → `last_score`) gives you this for free. ## Next Steps -- [RL Framework Integration](rl-integration.md) - Use rewards in training -- [Environment Anatomy](environment-anatomy.md) - Where to implement rewards -::: +- [Rubrics tutorial](../tutorials/rubrics.md) — full API walkthrough with composable examples. +- [RFC 004](https://github.com/meta-pytorch/OpenEnv/blob/main/rfcs/004-rubrics.md) — design rationale. +- [RL Framework Integration](rl-integration.md) — consume the reward signal in a training loop. +- [Environment Anatomy](environment-anatomy.md) — where the rubric plugs in. diff --git a/docs/source/tutorials/index.md b/docs/source/tutorials/index.md index 0bc7be4d7a..1e877bfea4 100644 --- a/docs/source/tutorials/index.md +++ b/docs/source/tutorials/index.md @@ -9,6 +9,7 @@ If you're new to OpenEnv, we recommend starting with the [Getting Started](/auto ## Available Tutorials - **[OpenEnv Tutorial](openenv-tutorial.md)** - A comprehensive introduction to OpenEnv, covering installation, basic usage, and core concepts. +- **[Rubrics](rubrics.md)** - Compose reward functions from reusable pieces using `openenv.core.rubrics` — `Gate`, `WeightedSum`, `LLMJudge`, and `TrajectoryRubric`. - **[Wordle GRPO Training](wordle-grpo.md)** - Learn how to train an agent to play Wordle using Group Relative Policy Optimization (GRPO). - **[RL Training with 2048](rl-training-2048.md)** - Train a language model to play 2048 using GRPO reinforcement learning. *(GPU Required)* @@ -16,6 +17,7 @@ If you're new to OpenEnv, we recommend starting with the [Getting Started](/auto :maxdepth: 2 :hidden: openenv-tutorial +rubrics wordle-grpo rl-training-2048 ``` diff --git a/docs/source/tutorials/rubrics.md b/docs/source/tutorials/rubrics.md new file mode 100644 index 0000000000..f7b1f45cc3 --- /dev/null +++ b/docs/source/tutorials/rubrics.md @@ -0,0 +1,354 @@ +# Rubrics: Composable Reward Computation + +Rubrics are OpenEnv's first-class abstraction for computing rewards. They let you build multi-criteria reward functions from small reusable pieces. This tutorial walks through the API end-to-end, from a one-line rubric to a full environment that introspects its reward signal at training time. + +## Why Rubrics? + +Before rubrics, each environment rolled its own reward logic. Three pain points surfaced repeatedly: + +1. **No standard interface**. Every environment author invented their own `compute_reward(...)` shape, so reusing a reward component across environments meant copy-pasting. +2. **Multi-criteria evaluation was ad-hoc**. "Code must compile, tests must pass, style matters a bit" becomes a tangle of nested `if`/`else` and hand-rolled weighted averages. There was no consistent way to ask *which* criterion caused a low reward. +3. **LLM judges and sandboxed checks are slow**. Without a framework-level concept of "reward component", batch evaluation couldn't parallelise the I/O-bound pieces. + +The Rubric API is small: you subclass, implement `forward`, and the framework gives you composition, introspection, and parallel evaluation for free. + +## Your First Rubric + +A rubric is a callable with a `forward(action, observation) -> float` method. + +```python +from openenv.core.rubrics import Rubric + + +class MessageLengthRubric(Rubric): + """Reward 1.0 if the message is 5–20 characters long, else 0.0.""" + + def forward(self, action, observation) -> float: + length = len(action.message) + return 1.0 if 5 <= length <= 20 else 0.0 +``` + +That's the whole contract. Instantiate it and call it: + +```python +rubric = MessageLengthRubric() +score = rubric(action, observation) # runs forward + hooks +print(rubric.last_score) # latest score is cached on the rubric +``` + +`Rubric.__call__` runs pre- and post-hooks around your `forward`, caches the result on `self.last_score`, and supports async `forward` implementations transparently. (If you've used `torch.nn.Module`, the subclass-and-implement-`forward` pattern will feel familiar — children assigned as instance attributes auto-register with the parent.) + +### Optional hooks for observability + +You can attach hooks without subclassing — useful for logging every component's score without polluting `forward`: + +```python +def log_score(rubric, action, obs, result): + print(f"{type(rubric).__name__}: {result:.2f}") + +rubric.register_forward_hook(log_score) +``` + +### State dict + +Rubrics implement `state_dict()` / `load_state_dict(state)` so their configuration (thresholds, prompt templates, etc.) can be serialised alongside model checkpoints. The default implementations return an empty dict — override them when your rubric has tunable parameters. + +## Composing Rubrics + +The real power shows up when you stack rubrics. `openenv.core.rubrics` ships with four containers. + +### `WeightedSum` — multi-criteria averaging + +Use when several independent criteria each contribute to the final score. + +```python +from openenv.core.rubrics import WeightedSum + + +class TestsPassRubric(Rubric): + def forward(self, action, observation) -> float: + return observation.tests_passed / max(observation.tests_total, 1) + + +class StyleRubric(Rubric): + def forward(self, action, observation) -> float: + return 1.0 if action.code.count("\n\n\n") == 0 else 0.6 + + +reward = WeightedSum( + [TestsPassRubric(), StyleRubric()], + weights=[0.7, 0.3], +) +``` + +Weights must sum to `1.0`. `WeightedSum` evaluates its children with `asyncio.gather` when any of them is async, so an LLM-backed child does not block the synchronous ones. + +### `Gate` — hard constraints + +Use when a child score below a threshold should short-circuit the reward to zero. + +```python +from openenv.core.rubrics import Gate + +reward = Gate(TestsPassRubric(), threshold=0.5) # 0.0 if fewer than half the tests pass +``` + +`Gate` returns `0.0` when the child score is below the threshold, and passes the child score through unchanged otherwise. + +### `Sequential` — fail-fast pipeline + +Use when criteria are ordered: a later criterion only matters if the earlier ones passed. Sequential returns `0.0` the moment any child returns `0.0` and does not evaluate the remaining children — great for gating expensive checks like sandboxed test runs or LLM calls. + +```python +from openenv.core.rubrics import Sequential + +reward = Sequential( + Gate(CompilesRubric(), threshold=1.0), # skip everything if it doesn't compile + Gate(TestsPassRubric(), threshold=0.5), # and skip style if tests are failing + WeightedSum([TestsPassRubric(), StyleRubric()], [0.7, 0.3]), +) +``` + +### `RubricList` and `RubricDict` — dynamic dispatch + +When the right rubric depends on the current observation (e.g. one rubric per game in a multi-game environment), wrap the options in a `RubricList` or `RubricDict` and dispatch in your parent rubric's `forward`. + +```python +from openenv.core.rubrics import Rubric, RubricDict + + +class MultiGameRubric(Rubric): + def __init__(self): + super().__init__() + self.games = RubricDict({ + "pong": PongRubric(), + "breakout": BreakoutRubric(), + }) + + def forward(self, action, observation) -> float: + return self.games[observation.game_id](action, observation) +``` + +`RubricList` and `RubricDict` do not aggregate on their own — calling them directly raises. Their job is auto-registration (so their children show up in `named_rubrics()`) and indexed access. + +### Introspection: `named_rubrics()` + +Assigning a child rubric as an attribute auto-registers it with the parent. Training code can then walk the tree: + +```python +composite = WeightedSum( + [Gate(CompilesRubric(), 1.0), TestsPassRubric(), StyleRubric()], + [0.2, 0.5, 0.3], +) + +for name, child in composite.named_rubrics(): + print(f"{name:30s} last_score={child.last_score}") +``` + +After running the composite once, every component's most recent score is cached on `last_score` — no manual bookkeeping. + +## LLM-as-judge: `LLMJudge` + +When a criterion is too subjective for a handwritten heuristic ("is this argument persuasive?", "is this explanation clear?"), use an LLM as the judge. `LLMJudge` wraps an `LLMClient` with a prompt template and a score extractor. + +Any OpenAI-compatible endpoint works: hosted OpenAI / Anthropic, or open-weight models served through vLLM, Ollama, Hugging Face Inference Providers, etc. Pick a client and hand it to `LLMJudge`: + +```python +import os + +from openenv.core.llm_client import OpenAIClient, create_llm_client +from openenv.core.rubrics import LLMJudge + +# Option 1 — hosted OpenAI (the factory also supports "anthropic"). +client = create_llm_client( + "openai", + model="gpt-4.1-mini", + api_key=os.environ["OPENAI_API_KEY"], +) + +# Option 2 — open-weight model served via a local OpenAI-compatible endpoint +# (vLLM, Ollama, Hugging Face Inference Providers, …). Point OpenAIClient +# at the base URL and the model id the server exposes. `api_key` is optional +# and defaults to "not-needed" for local endpoints. +client = OpenAIClient( + endpoint="http://localhost", + port=8000, + model="Qwen/Qwen3-1.7B", +) + +clarity_judge = LLMJudge( + client=client, + prompt_template=( + "Rate the clarity of this explanation on a 0-10 scale. " + "Reply with the number only.\n\n" + "Explanation:\n{action}\n" + ), + score_pattern=r"(\d+(?:\.\d+)?)", + normalize=True, # clamps extracted score to [0, 1] +) +``` + +`LLMJudge.forward` is async. When you put it inside `WeightedSum` or `Sequential`, the container awaits it transparently. A few caveats worth stating up front: + +- **Cost and latency** scale with the number of episodes and the number of rubric calls per step. `Sequential` + `Gate` earlier in the pipeline is the usual answer. +- **Determinism** is not free. Cache scores when you can, and consider temperature 0 for repeatable eval runs. +- **API keys** belong in environment variables (`OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, …), not in code that ships to the Hub. + +## Delayed Rewards: `TrajectoryRubric` + +Some signals only materialise at the end of an episode — chess win/loss, unit-test suite success, a goal reached after many steps. `TrajectoryRubric` accumulates `(action, observation)` pairs internally and only invokes your scoring logic on the terminal observation. + +```python +from openenv.core.rubrics import TrajectoryRubric + + +class WinLossRubric(TrajectoryRubric): + def score_trajectory(self, trajectory) -> float: + _, final_obs = trajectory[-1] + return final_obs.final_reward # +1 win, -1 loss, 0 draw + + def compute_step_rewards(self): + # Credit assignment: distribute the final score across steps however you like. + final = self.score_trajectory(self._trajectory) + return [final] * len(self._trajectory) +``` + +`forward(action, obs)` returns `intermediate_reward` (default `0.0`) until `observation.done` is `True`, then calls `score_trajectory`. After the episode ends, training code can pull per-step rewards via `rubric.compute_step_rewards()`. + +For the common exponentially-discounted case, subclass `ExponentialDiscountingTrajectoryRubric` instead and only implement `score_trajectory`: + +```python +from openenv.core.rubrics import ExponentialDiscountingTrajectoryRubric + + +class ChessOutcomeRubric(ExponentialDiscountingTrajectoryRubric): + def score_trajectory(self, trajectory) -> float: + _, final_obs = trajectory[-1] + return final_obs.reward # already +1 / 0 / -1 from the engine +``` + +This is exactly the pattern the built-in `envs/chess_env/` uses — see `envs/chess_env/server/rubrics.py` for the complete real-world example. + +:::{caution} +The `TrajectoryRubric` keeps the trajectory in CPU memory. If your observation carries GPU tensors (images, embeddings), detach and move them to CPU before returning from `step()` — otherwise the trajectory holds onto GPU memory across the whole episode. +::: + +## Wiring a Rubric into an `Environment` + +Rubrics are **server-side**. Each environment declares its rubric in `__init__`, and `step` runs it via the `_apply_rubric` helper. The base `Environment` class accepts the rubric through its constructor and stores it as `self.rubric`. + +Here is a complete minimal environment that composes a `Gate` with a `WeightedSum` and exposes the reward through its observation: + +```python +from openenv.core.env_server.interfaces import Environment +from openenv.core.env_server.types import Action, Observation, State +from openenv.core.rubrics import Gate, Rubric, WeightedSum + + +class CodeAction(Action): + code: str + + +class CodeObservation(Observation): + compiles: bool = False + tests_passed: int = 0 + tests_total: int = 0 + + +class CodeState(State): + attempts: int = 0 + + +class CompilesRubric(Rubric): + def forward(self, action, observation) -> float: + return 1.0 if observation.compiles else 0.0 + + +class TestsPassRubric(Rubric): + def forward(self, action, observation) -> float: + if observation.tests_total == 0: + return 0.0 + return observation.tests_passed / observation.tests_total + + +class StyleRubric(Rubric): + def forward(self, action, observation) -> float: + return 1.0 if action.code.count("\n\n\n") == 0 else 0.6 + + +def build_code_rubric() -> Rubric: + return WeightedSum( + [ + Gate(CompilesRubric(), threshold=1.0), # gate everything on compilation + TestsPassRubric(), + StyleRubric(), + ], + weights=[0.1, 0.7, 0.2], + ) + + +class CodeEnvironment(Environment[CodeAction, CodeObservation, CodeState]): + def __init__(self): + super().__init__(rubric=build_code_rubric()) + self._state = CodeState() + + def reset(self, seed=None, episode_id=None, **kwargs) -> CodeObservation: + self._reset_rubric() # clear any trajectory / cached last_score + self._state = CodeState() + return CodeObservation() + + def step(self, action: CodeAction, timeout_s=None, **kwargs) -> CodeObservation: + self._state.attempts += 1 + obs = self._run_code(action) # your domain-specific execution + obs.reward = self._apply_rubric(action, obs) + return obs + + @property + def state(self) -> CodeState: + return self._state + + def _run_code(self, action: CodeAction) -> CodeObservation: + # Placeholder for whatever your environment actually does. + compiles = "def " in action.code + return CodeObservation( + compiles=compiles, + tests_passed=3 if compiles else 0, + tests_total=3, + ) +``` + +The three pieces the base class expects from you: + +1. **Pass the rubric to `super().__init__(rubric=...)`** so `self.rubric` is set. +2. **Call `self._reset_rubric()` from `reset`** so trajectory state does not leak between episodes. +3. **Call `self._apply_rubric(action, obs)` from `step`** and attach the result to `obs.reward`. There is also `_apply_rubric_async` for `step_async`. + +### Inspecting rewards from training code + +Because children are auto-registered, the training loop can walk the rubric tree and log component-level diagnostics without the environment exposing a custom API: + +```python +env = CodeEnvironment() +obs = env.reset() +obs = env.step(CodeAction(code="def solution(): return 42")) + +for name, component in env.rubric.named_rubrics(): + print(f"{name:30s} last_score={component.last_score:.2f}") +``` + +That snippet works for *any* OpenEnv environment that sets `self.rubric`, regardless of whether the rubric is a single scalar or a deeply nested composition. + +### Where the reward ends up during training + +Training frameworks consume the reward through the same channel as any other OpenEnv observation field: `step()` returns an `Observation` whose `reward` is the rubric's output, and the client delivers it via `result.reward`. + +With [TRL](https://huggingface.co/docs/trl/main/en/openenv), the recommended path is `GRPOTrainer`'s `environment_factory`: you define a thin wrapper class with tool methods that call the OpenEnv client, store `self.reward = result.observation.reward` after each step, and a plain reward function reads it off the `environments` parameter. The [TRL OpenEnv integration guide](https://huggingface.co/docs/trl/main/en/openenv) has the full recipe, and [`examples/scripts/openenv/`](https://github.com/huggingface/trl/tree/main/examples/scripts/openenv) ships ready-to-run scripts. The same observation shape works with [torchforge](https://github.com/pytorch-labs/torchforge) and other OpenEnv-compatible training stacks. + +`named_rubrics()` is orthogonal: use it to **log per-component scores** (to Weights & Biases, TensorBoard, trackio, …) while training, without changing the reward the optimiser sees. + +## Next Steps + +- **Real-world trajectory example** — walk through `envs/chess_env/server/rubrics.py` and `chess_environment.py` to see `ExponentialDiscountingTrajectoryRubric` wired into a game environment. +- **Design details** — [RFC 004](https://github.com/meta-pytorch/OpenEnv/blob/main/rfcs/004-rubrics.md) covers the rationale for the composable API and the "rewards inside the environment" invariant. +- **Reward design basics** — the [Reward Design](../guides/rewards.md) guide covers sparse-vs-dense signals and common pitfalls that still apply on top of any rubric composition. +- **Training loop integration** — see the [RL Framework Integration](../guides/rl-integration.md) guide and the [TRL OpenEnv integration guide](https://huggingface.co/docs/trl/main/en/openenv) for the recommended `environment_factory` pattern. From 6c2367989a0ad636ec2f71f04ec6952b18fcb217 Mon Sep 17 00:00:00 2001 From: Sergio Paniego Blanco Date: Mon, 20 Apr 2026 17:37:57 +0200 Subject: [PATCH 2/6] Update docs/source/tutorials/rubrics.md Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- docs/source/tutorials/rubrics.md | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/docs/source/tutorials/rubrics.md b/docs/source/tutorials/rubrics.md index f7b1f45cc3..64bf5498bb 100644 --- a/docs/source/tutorials/rubrics.md +++ b/docs/source/tutorials/rubrics.md @@ -277,13 +277,16 @@ class StyleRubric(Rubric): def build_code_rubric() -> Rubric: - return WeightedSum( - [ - Gate(CompilesRubric(), threshold=1.0), # gate everything on compilation - TestsPassRubric(), - StyleRubric(), - ], - weights=[0.1, 0.7, 0.2], +def build_code_rubric() -> Rubric: + return Sequential( + Gate(CompilesRubric(), threshold=1.0), # gate everything on compilation + WeightedSum( + [ + TestsPassRubric(), + StyleRubric(), + ], + weights=[0.7, 0.3], + ), ) From cb47161de1197ccb8d2a12e2a580ea7113c9a676 Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Mon, 20 Apr 2026 17:40:17 +0200 Subject: [PATCH 3/6] Follow-up to review suggestion: dedupe def line, add Sequential import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The GitHub "commit suggestion" that switched `build_code_rubric` to `Sequential(Gate(...), WeightedSum(...))` landed with two rough edges: 1. A duplicate `def build_code_rubric() -> Rubric:` line at the top of the function. 2. The `CodeEnvironment` code block imports only `Gate, Rubric, WeightedSum` — `Sequential` was missing now that the composition uses it. Also update the intro prose from "composes a `Gate` with a `WeightedSum`" to "composes a `Sequential` gate-then-`WeightedSum` pipeline" to match the actual shape. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/source/tutorials/rubrics.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/source/tutorials/rubrics.md b/docs/source/tutorials/rubrics.md index 64bf5498bb..7873d47241 100644 --- a/docs/source/tutorials/rubrics.md +++ b/docs/source/tutorials/rubrics.md @@ -237,12 +237,12 @@ The `TrajectoryRubric` keeps the trajectory in CPU memory. If your observation c Rubrics are **server-side**. Each environment declares its rubric in `__init__`, and `step` runs it via the `_apply_rubric` helper. The base `Environment` class accepts the rubric through its constructor and stores it as `self.rubric`. -Here is a complete minimal environment that composes a `Gate` with a `WeightedSum` and exposes the reward through its observation: +Here is a complete minimal environment that composes a `Sequential` gate-then-`WeightedSum` pipeline and exposes the reward through its observation: ```python from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import Action, Observation, State -from openenv.core.rubrics import Gate, Rubric, WeightedSum +from openenv.core.rubrics import Gate, Rubric, Sequential, WeightedSum class CodeAction(Action): @@ -276,7 +276,6 @@ class StyleRubric(Rubric): return 1.0 if action.code.count("\n\n\n") == 0 else 0.6 -def build_code_rubric() -> Rubric: def build_code_rubric() -> Rubric: return Sequential( Gate(CompilesRubric(), threshold=1.0), # gate everything on compilation From a92cfec5a8bc2127944fdb28042b02c9734dee93 Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Mon, 20 Apr 2026 17:43:06 +0200 Subject: [PATCH 4/6] Address deep-review feedback Findings from a pre-merge review of the rubrics tutorial: - `WinLossRubric.score_trajectory` returned `final_obs.final_reward`, but `Observation` only exposes `.reward`. Chess env itself uses `final_obs.reward`. Fix the attribute name so the sample runs. - The hooks paragraph asserted "pre- and post-hooks", but the code sample only showed `register_forward_hook` (post) and gave no hint about when either fires. Add `register_forward_pre_hook` and a one-sentence clarifier about firing order + async awaiting. - `RubricList` / `RubricDict` described *how* (no aggregation, raises on direct call) but not *why* (dynamic membership / observation- dependent dispatch). Add a short rationale so readers can tell them apart from plain attribute children. - `TrajectoryRubric` silently holds on to trajectory memory if `observation.done` is never `True`. That's a real footgun for anyone wiring up a new environment. Add a caution box. - `concepts.md` used "delayed trajectory-based scoring" while the tutorial section is called "Delayed Rewards: TrajectoryRubric". Unify on "delayed rewards" and reference `TrajectoryRubric` explicitly. - The `environment-anatomy.md` rubric section restated the 3-step contract already covered in the tutorial's wiring section. Tighten to a single paragraph that points at the tutorial for the full API. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/source/concepts.md | 2 +- docs/source/guides/environment-anatomy.md | 8 +------- docs/source/tutorials/rubrics.md | 13 +++++++++---- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/docs/source/concepts.md b/docs/source/concepts.md index 3a803bc542..a3dd739e7e 100644 --- a/docs/source/concepts.md +++ b/docs/source/concepts.md @@ -56,7 +56,7 @@ A **StepResult** bundles together everything returned from a step: ### Rubric -A **Rubric** is a composable unit of reward computation that lives inside the environment. Rubrics can be combined (`WeightedSum`, `Gate`, `Sequential`) to express multi-criteria rewards, support LLM judges for subjective criteria, and support delayed trajectory-based scoring. Every environment declares `self.rubric` in its constructor; `step` calls it via `self._apply_rubric(action, observation)` and the result flows back through `observation.reward`. See the [Rubrics tutorial](tutorials/rubrics.md) for the full API. +A **Rubric** is a composable unit of reward computation that lives inside the environment. Rubrics can be combined (`WeightedSum`, `Gate`, `Sequential`) to express multi-criteria rewards, support LLM judges for subjective criteria, and handle delayed rewards via `TrajectoryRubric`. Every environment declares `self.rubric` in its constructor; `step` calls it via `self._apply_rubric(action, observation)` and the result flows back through `observation.reward`. See the [Rubrics tutorial](tutorials/rubrics.md) for the full API. ### Client diff --git a/docs/source/guides/environment-anatomy.md b/docs/source/guides/environment-anatomy.md index 647ed88cac..bafeae7302 100644 --- a/docs/source/guides/environment-anatomy.md +++ b/docs/source/guides/environment-anatomy.md @@ -92,13 +92,7 @@ app = create_server(MyEnvironment) ## Rewards via the Rubric -Rewards in OpenEnv are computed **inside the environment** — they never come from outside the `step` call. The base `Environment` accepts an optional `rubric` parameter on `__init__` and stores it as `self.rubric`. The recommended flow is: - -1. Build the rubric and pass it to `super().__init__(rubric=...)` so `self.rubric` is set. -2. Call `self._reset_rubric()` at the start of `reset` so trajectory state does not leak between episodes. -3. Call `self._apply_rubric(action, observation)` inside `step` and attach the result to `observation.reward`. Use `self._apply_rubric_async(...)` from `step_async`. - -The rubric system (`openenv.core.rubrics`) ships with composable containers (`WeightedSum`, `Gate`, `Sequential`), an `LLMJudge` for subjective criteria, and `TrajectoryRubric` for delayed rewards. The [Rubrics tutorial](../tutorials/rubrics.md) covers the full API. +Rewards are computed **inside the environment**, not by external code. The base `Environment` accepts an optional `rubric` on `__init__` — pass it to `super().__init__(rubric=...)`, call `self._reset_rubric()` from `reset`, and `self._apply_rubric(action, observation)` from `step` (or `_apply_rubric_async` from `step_async`). The [Rubrics tutorial](../tutorials/rubrics.md) covers the composable API end-to-end. ## Next Steps diff --git a/docs/source/tutorials/rubrics.md b/docs/source/tutorials/rubrics.md index 7873d47241..f6561a8e46 100644 --- a/docs/source/tutorials/rubrics.md +++ b/docs/source/tutorials/rubrics.md @@ -40,13 +40,14 @@ print(rubric.last_score) # latest score is cached on the rubric ### Optional hooks for observability -You can attach hooks without subclassing — useful for logging every component's score without polluting `forward`: +You can attach hooks without subclassing — useful for logging every component's score without polluting `forward`. Post-hooks run after `forward` completes and see the returned score; pre-hooks run before `forward` and are handy for input validation or instrumentation. When a rubric is async, hooks are awaited transparently. ```python def log_score(rubric, action, obs, result): print(f"{type(rubric).__name__}: {result:.2f}") -rubric.register_forward_hook(log_score) +rubric.register_forward_hook(log_score) # fires after forward() +rubric.register_forward_pre_hook(lambda r, a, o: None) # fires before forward() ``` ### State dict @@ -129,7 +130,7 @@ class MultiGameRubric(Rubric): return self.games[observation.game_id](action, observation) ``` -`RubricList` and `RubricDict` do not aggregate on their own — calling them directly raises. Their job is auto-registration (so their children show up in `named_rubrics()`) and indexed access. +`RubricList` and `RubricDict` do not aggregate on their own — calling them directly raises. Their job is auto-registration (so their children show up in `named_rubrics()`) and indexed access. Reach for them when the parent rubric needs to pick a child *at runtime* based on the observation — if the set of children is fixed, plain attributes are simpler. ### Introspection: `named_rubrics()` @@ -205,7 +206,7 @@ from openenv.core.rubrics import TrajectoryRubric class WinLossRubric(TrajectoryRubric): def score_trajectory(self, trajectory) -> float: _, final_obs = trajectory[-1] - return final_obs.final_reward # +1 win, -1 loss, 0 draw + return final_obs.reward # +1 win, -1 loss, 0 draw def compute_step_rewards(self): # Credit assignment: distribute the final score across steps however you like. @@ -215,6 +216,10 @@ class WinLossRubric(TrajectoryRubric): `forward(action, obs)` returns `intermediate_reward` (default `0.0`) until `observation.done` is `True`, then calls `score_trajectory`. After the episode ends, training code can pull per-step rewards via `rubric.compute_step_rewards()`. +:::{caution} +If `observation.done` never becomes `True`, `score_trajectory` is never called and the trajectory grows unbounded in memory. Make sure `step` flips `done` on every terminal transition, and call `self._reset_rubric()` in `Environment.reset` so trajectories do not leak across episodes. +::: + For the common exponentially-discounted case, subclass `ExponentialDiscountingTrajectoryRubric` instead and only implement `score_trajectory`: ```python From d181c446fa6b4fe803880c6b722aec65d34885a5 Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Tue, 21 Apr 2026 12:39:31 +0200 Subject: [PATCH 5/6] docs(rubrics): document hybrid wiring, expand credit-assignment, add eval section MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three additions from a deep pass on the tutorial, motivated by reviewer feedback and a broader audit: - The 3-point environment wiring contract assumed "pure rubric" mode (obs.reward = self._apply_rubric(...)), but envs/chess_env/ and envs/carla_env/ deliberately don't follow it — obs.reward comes from game mechanics / handcrafted signals and the rubric is only called for trajectory accumulation. Add a {note} under the contract calling the hybrid pattern out explicitly and pointing at both envs. - Expand the compute_step_rewards paragraph on TrajectoryRubric to explain what training code actually does with the per-step list (credit assignment for advantage estimation / return-to-go) and when to override vs use ExponentialDiscountingTrajectoryRubric directly. - Add a "Using Rubrics for Evaluation" section showing rubrics as plain callables over a static dataset, with per-component breakdown via named_rubrics() / last_score. Deliberately does not claim integration with src/openenv/core/evals/ — that bridge does not exist on main today and is tracked as a follow-up for when PR #471 / feature/harness-collect land. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/source/tutorials/rubrics.md | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/docs/source/tutorials/rubrics.md b/docs/source/tutorials/rubrics.md index f6561a8e46..dd0044fb4a 100644 --- a/docs/source/tutorials/rubrics.md +++ b/docs/source/tutorials/rubrics.md @@ -214,7 +214,7 @@ class WinLossRubric(TrajectoryRubric): return [final] * len(self._trajectory) ``` -`forward(action, obs)` returns `intermediate_reward` (default `0.0`) until `observation.done` is `True`, then calls `score_trajectory`. After the episode ends, training code can pull per-step rewards via `rubric.compute_step_rewards()`. +`forward(action, obs)` returns `intermediate_reward` (default `0.0`) until `observation.done` is `True`, then calls `score_trajectory`. After the episode ends, call `rubric.compute_step_rewards()` to get one reward per step — same length as the trajectory. This is the hook for credit assignment: training code feeds these per-step rewards back into advantage estimation, return-to-go, or whatever your optimizer expects. `ExponentialDiscountingTrajectoryRubric` precomputes `gamma^(T-1-t) * final_score` for you; override `compute_step_rewards` in your subclass if you want a different strategy (all-to-last, equal split, task-specific shaping). :::{caution} If `observation.done` never becomes `True`, `score_trajectory` is never called and the trajectory grows unbounded in memory. Make sure `step` flips `done` on every terminal transition, and call `self._reset_rubric()` in `Environment.reset` so trajectories do not leak across episodes. @@ -330,6 +330,10 @@ The three pieces the base class expects from you: 2. **Call `self._reset_rubric()` from `reset`** so trajectory state does not leak between episodes. 3. **Call `self._apply_rubric(action, obs)` from `step`** and attach the result to `obs.reward`. There is also `_apply_rubric_async` for `step_async`. +:::{note} +Some environments already compute `obs.reward` from game mechanics or a handcrafted multi-component signal (see `envs/chess_env/` and `envs/carla_env/`). In that case, call `self._apply_rubric(action, obs)` without assigning its return value — the rubric still accumulates the trajectory for `compute_step_rewards()` and still exposes per-component scores via `named_rubrics()`, but `obs.reward` stays authoritative. +::: + ### Inspecting rewards from training code Because children are auto-registered, the training loop can walk the rubric tree and log component-level diagnostics without the environment exposing a custom API: @@ -353,6 +357,24 @@ With [TRL](https://huggingface.co/docs/trl/main/en/openenv), the recommended pat `named_rubrics()` is orthogonal: use it to **log per-component scores** (to Weights & Biases, TensorBoard, trackio, …) while training, without changing the reward the optimiser sees. +## Using Rubrics for Evaluation + +A rubric is just a callable — nothing forces you to run it inside a training loop. Drop it into a for-loop over a static dataset and you have a multi-criteria scoring function for offline eval: + +```python +rubric = build_code_rubric() + +scores = [] +for action, obs in eval_dataset: + scores.append(rubric(action, obs)) + +print(f"mean reward: {sum(scores) / len(scores):.3f}") +for name, component in rubric.named_rubrics(): + print(f" {name:30s} last_score={component.last_score:.3f}") +``` + +The same rubric object used to compute training rewards doubles as the eval metric — one source of truth for "what is a good response". Per-component `last_score` gives you a per-criterion breakdown for free (useful for regression dashboards and failure analysis). When a component like `LLMJudge` is async, wrap the loop with `asyncio.run(...)` and `await rubric(action, obs)` so the judge calls can overlap. + ## Next Steps - **Real-world trajectory example** — walk through `envs/chess_env/server/rubrics.py` and `chess_environment.py` to see `ExponentialDiscountingTrajectoryRubric` wired into a game environment. From b203d4315033e52b1fcefdabc1b7fc1b5b02df68 Mon Sep 17 00:00:00 2001 From: sergiopaniego Date: Tue, 5 May 2026 11:18:43 +0200 Subject: [PATCH 6/6] docs(rubrics): add companion Colab notebook and badge Co-Authored-By: Claude Sonnet 4.6 --- docs/source/tutorials/rubrics.md | 2 + examples/rubrics.ipynb | 460 +++++++++++++++++++++++++++++++ 2 files changed, 462 insertions(+) create mode 100644 examples/rubrics.ipynb diff --git a/docs/source/tutorials/rubrics.md b/docs/source/tutorials/rubrics.md index dd0044fb4a..77abaafb50 100644 --- a/docs/source/tutorials/rubrics.md +++ b/docs/source/tutorials/rubrics.md @@ -1,5 +1,7 @@ # Rubrics: Composable Reward Computation +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/meta-pytorch/OpenEnv/blob/main/examples/rubrics.ipynb) + Rubrics are OpenEnv's first-class abstraction for computing rewards. They let you build multi-criteria reward functions from small reusable pieces. This tutorial walks through the API end-to-end, from a one-line rubric to a full environment that introspects its reward signal at training time. ## Why Rubrics? diff --git a/examples/rubrics.ipynb b/examples/rubrics.ipynb new file mode 100644 index 0000000000..368d4ced19 --- /dev/null +++ b/examples/rubrics.ipynb @@ -0,0 +1,460 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1191c64d-62d2-47a8-b13f-e117d8466581", + "metadata": {}, + "source": "# Rubrics: Composable Reward Computation\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/meta-pytorch/OpenEnv/blob/main/examples/rubrics.ipynb)\n\nRubrics are OpenEnv's first-class abstraction for computing rewards. They let you build multi-criteria reward functions from small reusable pieces. This tutorial walks through the API end-to-end, from a one-line rubric to a full environment that introspects its reward signal at training time." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dab20999-1ef6-4902-8ca9-580d00262787", + "metadata": {}, + "outputs": [], + "source": "!pip install -q \"openenv-core @ git+https://github.com/meta-pytorch/OpenEnv.git\"" + }, + { + "cell_type": "markdown", + "id": "45392f8d-9746-4b30-bbc7-37827e888f08", + "metadata": {}, + "source": "## Why Rubrics?\n\nBefore rubrics, each environment rolled its own reward logic. Three pain points surfaced repeatedly:\n\n1. **No standard interface**. Every environment author invented their own `compute_reward(...)` shape, so reusing a reward component across environments meant copy-pasting.\n2. **Multi-criteria evaluation was ad-hoc**. \"Code must compile, tests must pass, style matters a bit\" becomes a tangle of nested `if`/`else` and hand-rolled weighted averages. There was no consistent way to ask *which* criterion caused a low reward.\n3. **LLM judges and sandboxed checks are slow**. Without a framework-level concept of \"reward component\", batch evaluation couldn't parallelise the I/O-bound pieces.\n\nThe Rubric API is small: you subclass, implement `forward`, and the framework gives you composition, introspection, and parallel evaluation for free." + }, + { + "cell_type": "markdown", + "id": "8c1cdf16-9f4c-4aa3-a8ef-8dd4f5da0557", + "metadata": {}, + "source": "## Your First Rubric\n\nA rubric is a callable with a `forward(action, observation) -> float` method." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29655a69-b32f-499e-b4d7-5d743f7098e7", + "metadata": {}, + "outputs": [], + "source": "from openenv.core.rubrics import Rubric\nfrom types import SimpleNamespace\n\n\nclass MessageLengthRubric(Rubric):\n \"\"\"Reward 1.0 if the message is 5\u201320 characters long, else 0.0.\"\"\"\n\n def forward(self, action, observation) -> float:\n length = len(action.message)\n return 1.0 if 5 <= length <= 20 else 0.0\n\n\nrubric = MessageLengthRubric()\nobs = SimpleNamespace()\n\naction = SimpleNamespace(message=\"hi\")\nscore = rubric(action, obs)\nprint(f\"'hi' \u2192 {score} (last_score={rubric.last_score})\")\n\naction = SimpleNamespace(message=\"hello world\")\nscore = rubric(action, obs)\nprint(f\"'hello world' \u2192 {score} (last_score={rubric.last_score})\")\n\naction = SimpleNamespace(message=\"this message is way too long for the rubric\")\nscore = rubric(action, obs)\nprint(f\"long message \u2192 {score} (last_score={rubric.last_score})\")" + }, + { + "cell_type": "markdown", + "id": "f6d6ecdb-fd6a-45cc-acd3-06ea0e2962a4", + "metadata": {}, + "source": "`Rubric.__call__` runs pre- and post-hooks around your `forward`, caches the result on `self.last_score`, and supports async `forward` implementations transparently. (If you've used `torch.nn.Module`, the subclass-and-implement-`forward` pattern will feel familiar \u2014 children assigned as instance attributes auto-register with the parent.)" + }, + { + "cell_type": "markdown", + "id": "bd54139b-3ab3-4596-be53-100b69a01be7", + "metadata": {}, + "source": "### Optional hooks for observability\n\nYou can attach hooks without subclassing \u2014 useful for logging every component's score without polluting `forward`. Post-hooks run after `forward` completes and see the returned score; pre-hooks run before `forward` and are handy for input validation or instrumentation. When a rubric is async, hooks are awaited transparently." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5948f2e6-02c9-416f-94e1-4712033e250e", + "metadata": {}, + "outputs": [], + "source": "def log_score(rubric, action, obs, result):\n print(f\"{type(rubric).__name__}: {result:.2f}\")\n\nrubric.register_forward_hook(log_score) # fires after forward()\nrubric.register_forward_pre_hook(lambda r, a, o: None) # fires before forward()\n\naction = SimpleNamespace(message=\"hello world\")\n_ = rubric(action, obs) # hook fires and prints" + }, + { + "cell_type": "markdown", + "id": "72b9e891-1fc8-41ec-84de-9d18baa0bce1", + "metadata": {}, + "source": "### State dict\n\nRubrics implement `state_dict()` / `load_state_dict(state)` so their configuration (thresholds, prompt templates, etc.) can be serialised alongside model checkpoints. The default implementations return an empty dict \u2014 override them when your rubric has tunable parameters." + }, + { + "cell_type": "markdown", + "id": "fa1f1cf4-06a3-4e26-a35c-69ab5d49d46d", + "metadata": {}, + "source": "## Composing Rubrics\n\nThe real power shows up when you stack rubrics. `openenv.core.rubrics` ships with four containers.\n\n### `WeightedSum` \u2014 multi-criteria averaging\n\nUse when several independent criteria each contribute to the final score." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "75b33b68-7741-41fd-a5c8-71beb462a0cf", + "metadata": {}, + "outputs": [], + "source": "from openenv.core.rubrics import WeightedSum\n\n\nclass TestsPassRubric(Rubric):\n def forward(self, action, observation) -> float:\n return observation.tests_passed / max(observation.tests_total, 1)\n\n\nclass StyleRubric(Rubric):\n def forward(self, action, observation) -> float:\n return 1.0 if action.code.count(\"\\n\\n\\n\") == 0 else 0.6\n\n\nreward = WeightedSum(\n [TestsPassRubric(), StyleRubric()],\n weights=[0.7, 0.3],\n)\n\naction = SimpleNamespace(code=\"def solution():\\n return 42\")\nobs = SimpleNamespace(tests_passed=3, tests_total=3)\nprint(f\"reward: {reward(action, obs):.2f}\")" + }, + { + "cell_type": "markdown", + "id": "de946621-6afa-4b1b-9c1a-e03741ca51c3", + "metadata": {}, + "source": "Weights must sum to `1.0`. `WeightedSum` evaluates its children with `asyncio.gather` when any of them is async, so an LLM-backed child does not block the synchronous ones." + }, + { + "cell_type": "markdown", + "id": "65b24bc4-3c41-47bc-b3b7-cf717491c893", + "metadata": {}, + "source": "### `Gate` \u2014 hard constraints\n\nUse when a child score below a threshold should short-circuit the reward to zero." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11dbb6e1-d4b6-41ea-903d-e559a125f9c2", + "metadata": {}, + "outputs": [], + "source": "from openenv.core.rubrics import Gate\n\nreward = Gate(TestsPassRubric(), threshold=0.5) # 0.0 if fewer than half the tests pass\n\nobs_fail = SimpleNamespace(tests_passed=1, tests_total=3)\nobs_pass = SimpleNamespace(tests_passed=2, tests_total=3)\naction = SimpleNamespace(code=\"x = 1\")\n\nprint(f\"1/3 tests pass \u2192 {reward(action, obs_fail):.2f}\") # below threshold \u2192 0.0\nprint(f\"2/3 tests pass \u2192 {reward(action, obs_pass):.2f}\") # above threshold \u2192 passes through" + }, + { + "cell_type": "markdown", + "id": "32e3a99b-811c-4b9b-b502-a5ea075f29e9", + "metadata": {}, + "source": "`Gate` returns `0.0` when the child score is below the threshold, and passes the child score through unchanged otherwise." + }, + { + "cell_type": "markdown", + "id": "13708f50-613d-4ee1-8d41-13811ad10cfc", + "metadata": {}, + "source": "### `Sequential` \u2014 fail-fast pipeline\n\nUse when criteria are ordered: a later criterion only matters if the earlier ones passed. Sequential returns `0.0` the moment any child returns `0.0` and does not evaluate the remaining children \u2014 great for gating expensive checks like sandboxed test runs or LLM calls." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f4767aa-092e-4da0-b0b4-e242ba16ed13", + "metadata": {}, + "outputs": [], + "source": [ + "import inspect\n", + "from openenv.core.rubrics import Sequential\n", + "\n", + "\n", + "class CompilesRubric(Rubric):\n", + " def forward(self, action, observation) -> float:\n", + " return 1.0 if observation.compiles else 0.0\n", + "\n", + "\n", + "reward = Sequential(\n", + " Gate(CompilesRubric(), threshold=1.0), # skip everything if it doesn't compile\n", + " Gate(TestsPassRubric(), threshold=0.5), # skip style if tests are failing\n", + " WeightedSum([TestsPassRubric(), StyleRubric()], [0.7, 0.3]),\n", + ")\n", + "\n", + "# Jupyter runs inside asyncio; Sequential detects the running loop and may return\n", + "# a coroutine even when all children are sync. This helper handles both cases.\n", + "async def call(rubric, action, obs):\n", + " result = rubric(action, obs)\n", + " return (await result) if inspect.iscoroutine(result) else result\n", + "\n", + "\n", + "# Doesn't compile \u2192 Gate short-circuits, returns 0.0 immediately\n", + "obs_bad = SimpleNamespace(compiles=False, tests_passed=0, tests_total=3)\n", + "action_bad = SimpleNamespace(code=\"bad code\")\n", + "print(f\"no compile \u2192 {await call(reward, action_bad, obs_bad):.2f}\")\n", + "\n", + "# Compiles + passes \u2192 runs the full WeightedSum\n", + "obs_good = SimpleNamespace(compiles=True, tests_passed=3, tests_total=3)\n", + "action_good = SimpleNamespace(code=\"def solution(): return 42\")\n", + "print(f\"compiles + passes \u2192 {await call(reward, action_good, obs_good):.2f}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "56af7d67-5410-44ac-ba35-9224bc762265", + "metadata": {}, + "source": "### `RubricList` and `RubricDict` \u2014 dynamic dispatch\n\nWhen the right rubric depends on the current observation (e.g. one rubric per game in a multi-game environment), wrap the options in a `RubricList` or `RubricDict` and dispatch in your parent rubric's `forward`." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "777653a4-c3b6-4509-a4b1-f3f239036e50", + "metadata": {}, + "outputs": [], + "source": "from openenv.core.rubrics import Rubric, RubricDict\n\n\nclass PongRubric(Rubric):\n def forward(self, action, observation) -> float:\n return 1.0 if observation.scored else 0.0\n\n\nclass BreakoutRubric(Rubric):\n def forward(self, action, observation) -> float:\n return observation.bricks_hit / max(observation.total_bricks, 1)\n\n\nclass MultiGameRubric(Rubric):\n def __init__(self):\n super().__init__()\n self.games = RubricDict({\n \"pong\": PongRubric(),\n \"breakout\": BreakoutRubric(),\n })\n\n def forward(self, action, observation) -> float:\n return self.games[observation.game_id](action, observation)\n\n\nmulti = MultiGameRubric()\n\npong_obs = SimpleNamespace(game_id=\"pong\", scored=True)\nbreakout_obs = SimpleNamespace(game_id=\"breakout\", bricks_hit=7, total_bricks=10)\naction = SimpleNamespace()\n\nprint(f\"pong: {multi(action, pong_obs):.2f}\")\nprint(f\"breakout: {multi(action, breakout_obs):.2f}\")" + }, + { + "cell_type": "markdown", + "id": "6a032e35-a7e9-4f87-a99d-728f745c1dc2", + "metadata": {}, + "source": "`RubricList` and `RubricDict` do not aggregate on their own \u2014 calling them directly raises. Their job is auto-registration (so their children show up in `named_rubrics()`) and indexed access." + }, + { + "cell_type": "markdown", + "id": "22cd32c1-fb05-4118-9956-e15f9f4a4d5d", + "metadata": {}, + "source": "### Introspection: `named_rubrics()`\n\nAssigning a child rubric as an attribute auto-registers it with the parent. Training code can then walk the tree:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "139d2827-8a52-4bf4-95db-26d305d32d69", + "metadata": {}, + "outputs": [], + "source": "composite = WeightedSum(\n [Gate(CompilesRubric(), 1.0), TestsPassRubric(), StyleRubric()],\n [0.2, 0.5, 0.3],\n)\n\nobs = SimpleNamespace(compiles=True, tests_passed=3, tests_total=3)\naction = SimpleNamespace(code=\"def solution(): return 42\")\n_ = composite(action, obs)\n\nfor name, child in composite.named_rubrics():\n score = f\"{child.last_score:.2f}\" if child.last_score is not None else \"None\"\n print(f\"{name:40s} last_score={score}\")" + }, + { + "cell_type": "markdown", + "id": "82e10806-a9df-4598-9974-64cb2cccbfbe", + "metadata": {}, + "source": "After running the composite once, every component's most recent score is cached on `last_score` \u2014 no manual bookkeeping." + }, + { + "cell_type": "markdown", + "id": "f13abe4c-ac21-4b13-827a-81222e3e07ea", + "metadata": {}, + "source": "## LLM-as-judge: `LLMJudge`\n\nWhen a criterion is too subjective for a handwritten heuristic (\"is this argument persuasive?\", \"is this explanation clear?\"), use an LLM as the judge. `LLMJudge` wraps an `LLMClient` with a prompt template and a score extractor.\n\nAny OpenAI-compatible endpoint works: hosted OpenAI / Anthropic, or open-weight models served through vLLM, Ollama, Hugging Face Inference Providers, etc.\n\n> **Note:** this cell requires an `OPENAI_API_KEY` environment variable. It is skipped automatically if the key is absent." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a25a00d6-698c-432c-ac5a-f202ff4b8453", + "metadata": {}, + "outputs": [], + "source": [ + "from google.colab import userdata\n", + "import os\n", + "\n", + "os.environ[\"OPENAI_API_KEY\"] = userdata.get(\"OPENAI_API_KEY\")\n", + "\n", + "import os\n", + "\n", + "from openenv.core.llm_client import create_llm_client\n", + "from openenv.core.rubrics import LLMJudge\n", + "\n", + "if not os.environ.get(\"OPENAI_API_KEY\"):\n", + " print(\"OPENAI_API_KEY not set \u2014 skipping LLMJudge demo. Set the key and re-run this cell.\")\n", + "else:\n", + " client = create_llm_client(\n", + " \"openai\",\n", + " model=\"gpt-4.1-mini\",\n", + " api_key=os.environ[\"OPENAI_API_KEY\"],\n", + " )\n", + "\n", + " clarity_judge = LLMJudge(\n", + " client=client,\n", + " prompt_template=(\n", + " \"Rate the clarity of this explanation on a 0-10 scale. \"\n", + " \"Reply with the number only.\\n\\n\"\n", + " \"Explanation:\\n{action}\\n\"\n", + " ),\n", + " score_pattern=r\"(\\d+(?:\\.\\d+)?)\",\n", + " normalize=True, # clamps extracted score to [0, 1]\n", + " )\n", + " print(\"LLMJudge created. Combine it with WeightedSum or Sequential like any other rubric.\")\n", + " print(\"clarity_judge.forward is async \u2014 await it or place it inside a WeightedSum container.\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "88b09256-13c6-4323-addf-52a38ee1b019", + "metadata": {}, + "source": "`LLMJudge.forward` is async. When you put it inside `WeightedSum` or `Sequential`, the container awaits it transparently. A few caveats worth stating up front:\n\n- **Cost and latency** scale with the number of episodes and the number of rubric calls per step. `Sequential` + `Gate` earlier in the pipeline is the usual answer.\n- **Determinism** is not free. Cache scores when you can, and consider temperature 0 for repeatable eval runs.\n- **API keys** belong in environment variables (`OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, \u2026), not in code that ships to the Hub." + }, + { + "cell_type": "markdown", + "id": "dff20b11-626f-40b2-8dd6-4d9069e025c6", + "metadata": {}, + "source": "## Delayed Rewards: `TrajectoryRubric`\n\nSome signals only materialise at the end of an episode \u2014 chess win/loss, unit-test suite success, a goal reached after many steps. `TrajectoryRubric` accumulates `(action, observation)` pairs internally and only invokes your scoring logic on the terminal observation." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de0e4677-27ad-4cf2-95b2-40cb8f550a0b", + "metadata": {}, + "outputs": [], + "source": "from openenv.core.rubrics import TrajectoryRubric\n\n\nclass WinLossRubric(TrajectoryRubric):\n def score_trajectory(self, trajectory) -> float:\n _, final_obs = trajectory[-1]\n return final_obs.reward # +1 win, -1 loss, 0 draw\n\n def compute_step_rewards(self):\n # Credit assignment: distribute the final score across steps.\n final = self.score_trajectory(self._trajectory)\n return [final] * len(self._trajectory)\n\n\nrubric = WinLossRubric()\n\n# Simulate a 3-step episode; reward only materialises on the last step\nfor t, done, r in [(0, False, 0.0), (1, False, 0.0), (2, True, 1.0)]:\n action = SimpleNamespace()\n obs = SimpleNamespace(done=done, reward=r)\n intermediate = rubric(action, obs)\n print(f\"step {t} (done={done}): intermediate_reward={intermediate}\")\n\nprint(f\"per-step rewards: {rubric.compute_step_rewards()}\")" + }, + { + "cell_type": "markdown", + "id": "d92e8c7e-cfdd-4ea4-ace7-6786f41d4d0c", + "metadata": {}, + "source": "`forward(action, obs)` returns `intermediate_reward` (default `0.0`) until `observation.done` is `True`, then calls `score_trajectory`. After the episode ends, call `rubric.compute_step_rewards()` to get one reward per step \u2014 same length as the trajectory.\n\nFor the common exponentially-discounted case, subclass `ExponentialDiscountingTrajectoryRubric` instead and only implement `score_trajectory`:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ce71e88a-7583-4ddb-8094-ef97655f85d6", + "metadata": {}, + "outputs": [], + "source": "from openenv.core.rubrics import ExponentialDiscountingTrajectoryRubric\n\n\nclass ChessOutcomeRubric(ExponentialDiscountingTrajectoryRubric):\n def score_trajectory(self, trajectory) -> float:\n _, final_obs = trajectory[-1]\n return final_obs.reward # already +1 / 0 / -1 from the engine\n\n# This is exactly the pattern the built-in envs/chess_env/ uses.\nprint(\"ChessOutcomeRubric defined \u2014 see envs/chess_env/server/rubrics.py for the full real-world example.\")" + }, + { + "cell_type": "markdown", + "id": "761c5ec1-4bf4-4eef-a9f3-400d3b57ab25", + "metadata": {}, + "source": "## Wiring a Rubric into an `Environment`\n\nRubrics are **server-side**. Each environment declares its rubric in `__init__`, and `step` runs it via the `_apply_rubric` helper. The base `Environment` class accepts the rubric through its constructor and stores it as `self.rubric`.\n\nHere is a complete minimal environment that composes a `Sequential` gate-then-`WeightedSum` pipeline and exposes the reward through its observation:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d342323a-7d4f-4808-8007-601f64450e8c", + "metadata": {}, + "outputs": [], + "source": [ + "from openenv.core.env_server.interfaces import Environment\n", + "from openenv.core.env_server.types import Action, Observation, State\n", + "from openenv.core.rubrics import Gate, Rubric, Sequential, WeightedSum\n", + "\n", + "\n", + "class CodeAction(Action):\n", + " code: str\n", + "\n", + "\n", + "class CodeObservation(Observation):\n", + " compiles: bool = False\n", + " tests_passed: int = 0\n", + " tests_total: int = 0\n", + "\n", + "\n", + "class CodeState(State):\n", + " attempts: int = 0\n", + "\n", + "\n", + "class CompilesRubric(Rubric):\n", + " def forward(self, action, observation) -> float:\n", + " return 1.0 if observation.compiles else 0.0\n", + "\n", + "\n", + "class TestsPassRubric(Rubric):\n", + " def forward(self, action, observation) -> float:\n", + " if observation.tests_total == 0:\n", + " return 0.0\n", + " return observation.tests_passed / observation.tests_total\n", + "\n", + "\n", + "class StyleRubric(Rubric):\n", + " def forward(self, action, observation) -> float:\n", + " return 1.0 if action.code.count(\"\\n\\n\\n\") == 0 else 0.6\n", + "\n", + "\n", + "def build_code_rubric() -> Rubric:\n", + " return Sequential(\n", + " Gate(CompilesRubric(), threshold=1.0), # gate everything on compilation\n", + " WeightedSum(\n", + " [TestsPassRubric(), StyleRubric()],\n", + " weights=[0.7, 0.3],\n", + " ),\n", + " )\n", + "\n", + "\n", + "class CodeEnvironment(Environment[CodeAction, CodeObservation, CodeState]):\n", + " def __init__(self):\n", + " super().__init__(rubric=build_code_rubric())\n", + " self._state = CodeState()\n", + "\n", + " def reset(self, seed=None, episode_id=None, **kwargs) -> CodeObservation:\n", + " self._reset_rubric() # clear any trajectory / cached last_score\n", + " self._state = CodeState()\n", + " return CodeObservation()\n", + "\n", + " async def step(self, action: CodeAction, timeout_s=None, **kwargs) -> CodeObservation:\n", + " self._state.attempts += 1\n", + " obs = self._run_code(action) # your domain-specific execution\n", + " obs.reward = await self._apply_rubric_async(action, obs)\n", + " return obs\n", + "\n", + " @property\n", + " def state(self) -> CodeState:\n", + " return self._state\n", + "\n", + " def _run_code(self, action: CodeAction) -> CodeObservation:\n", + " compiles = \"def \" in action.code\n", + " return CodeObservation(\n", + " compiles=compiles,\n", + " tests_passed=3 if compiles else 0,\n", + " tests_total=3,\n", + " )\n" + ] + }, + { + "cell_type": "markdown", + "id": "ed40bcf1-1f35-4b3d-a1b7-31ad33c87bac", + "metadata": {}, + "source": "The three pieces the base class expects from you:\n\n1. **Pass the rubric to `super().__init__(rubric=...)`** so `self.rubric` is set.\n2. **Call `self._reset_rubric()` from `reset`** so trajectory state does not leak between episodes.\n3. **Call `self._apply_rubric(action, obs)` from `step`** and attach the result to `obs.reward`." + }, + { + "cell_type": "markdown", + "id": "346635fc-042d-4435-8db4-56a19d01d2e7", + "metadata": {}, + "source": "### Inspecting rewards from training code\n\nBecause children are auto-registered, the training loop can walk the rubric tree and log component-level diagnostics without the environment exposing a custom API:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c5ea683-3640-45bf-8ef2-26a3b1ce2b49", + "metadata": {}, + "outputs": [], + "source": [ + "env = CodeEnvironment()\n", + "obs = env.reset()\n", + "obs = await env.step(CodeAction(code=\"def solution(): return 42\"))\n", + "\n", + "print(f\"total reward: {obs.reward:.2f}\")\n", + "print()\n", + "for name, component in env.rubric.named_rubrics():\n", + " score = f\"{component.last_score:.2f}\" if component.last_score is not None else \"None\"\n", + " print(f\"{name:40s} last_score={score}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "33cb2ac3-2635-434f-9081-f1bae555ed6d", + "metadata": {}, + "source": "That snippet works for *any* OpenEnv environment that sets `self.rubric`, regardless of whether the rubric is a single scalar or a deeply nested composition.\n\n### Where the reward ends up during training\n\nTraining frameworks consume the reward through the same channel as any other OpenEnv observation field: `step()` returns an `Observation` whose `reward` is the rubric's output, and the client delivers it via `result.reward`.\n\nWith [TRL](https://huggingface.co/docs/trl/main/en/openenv), the recommended path is `GRPOTrainer`'s `environment_factory`: you define a thin wrapper class with tool methods that call the OpenEnv client, store `self.reward = result.observation.reward` after each step, and a plain reward function reads it off the `environments` parameter.\n\n`named_rubrics()` is orthogonal: use it to **log per-component scores** (to Weights & Biases, TensorBoard, trackio, \u2026) while training, without changing the reward the optimiser sees." + }, + { + "cell_type": "markdown", + "id": "2262c753-8999-4ae7-8f62-6a53795ebd80", + "metadata": {}, + "source": "## Using Rubrics for Evaluation\n\nA rubric is just a callable \u2014 nothing forces you to run it inside a training loop. Drop it into a for-loop over a static dataset and you have a multi-criteria scoring function for offline eval:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "85a5535d-a108-48dc-8061-fbfccc52bf6a", + "metadata": {}, + "outputs": [], + "source": [ + "import inspect\n", + "\n", + "rubric = build_code_rubric()\n", + "\n", + "eval_dataset = [\n", + " (CodeAction(code=\"def add(a, b): return a + b\"),\n", + " CodeObservation(compiles=True, tests_passed=3, tests_total=3)),\n", + " (CodeAction(code=\"x = 1 + 2\"),\n", + " CodeObservation(compiles=False, tests_passed=0, tests_total=3)),\n", + " (CodeAction(code=\"def mul(a, b): return a * b\"),\n", + " CodeObservation(compiles=True, tests_passed=2, tests_total=3)),\n", + "]\n", + "\n", + "scores = []\n", + "for action, obs in eval_dataset:\n", + " result = rubric(action, obs)\n", + " if inspect.iscoroutine(result):\n", + " result = await result\n", + " scores.append(result)\n", + "\n", + "print(f\"mean reward: {sum(scores) / len(scores):.3f}\")\n", + "print()\n", + "print(\"Per-component scores (last sample):\")\n", + "for name, component in rubric.named_rubrics():\n", + " score = f\"{component.last_score:.2f}\" if component.last_score is not None else \"None\"\n", + " print(f\" {name:40s} last_score={score}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "764687f0-3a1f-4e1c-8108-0e815774839b", + "metadata": {}, + "source": "The same rubric object used to compute training rewards doubles as the eval metric \u2014 one source of truth for \"what is a good response\". Per-component `last_score` gives you a per-criterion breakdown for free (useful for regression dashboards and failure analysis)." + }, + { + "cell_type": "markdown", + "id": "765896d7-09ea-4a31-a560-a362cd17faeb", + "metadata": {}, + "source": "## Next Steps\n\n- **Real-world trajectory example** \u2014 walk through `envs/chess_env/server/rubrics.py` and `chess_environment.py` to see `ExponentialDiscountingTrajectoryRubric` wired into a game environment.\n- **Design details** \u2014 [RFC 004](https://github.com/meta-pytorch/OpenEnv/blob/main/rfcs/004-rubrics.md) covers the rationale for the composable API and the \"rewards inside the environment\" invariant.\n- **Reward design basics** \u2014 the [Reward Design](https://meta-pytorch.org/OpenEnv/guides/rewards.html) guide covers sparse-vs-dense signals and common pitfalls that still apply on top of any rubric composition.\n- **Training loop integration** \u2014 see the [RL Framework Integration](https://meta-pytorch.org/OpenEnv/guides/rl-integration.html) guide and the [TRL OpenEnv integration guide](https://huggingface.co/docs/trl/main/en/openenv) for the recommended `environment_factory` pattern." + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file