diff --git a/examples/echo_world_model/README.md b/examples/echo_world_model/README.md index 3fb3c1d42..8015f1f8d 100644 --- a/examples/echo_world_model/README.md +++ b/examples/echo_world_model/README.md @@ -83,13 +83,14 @@ and the rollout filters exist. | [`rollout.py`](./rollout.py) | Oracle + policy rollouts → role-tagged trajectories. | | [`train_echo.py`](./train_echo.py) | The CPU demo above. | | [`test_echo.py`](./test_echo.py) | Unit tests for the loss + masks (`python -m pytest test_echo.py`). | -| [`backends/`](./backends/) | The same config for the real GPU run: [SkyRL](./backends/skyrl.md) (open reference), [TRL](./backends/trl.md) (OpenEnv's recommended framework, with a runnable CPU [`trl_echo_demo.py`](./backends/trl_echo_demo.py)), [Tinker](./backends/tinker.md), [Foundry Fine-Tuning](./backends/foundry-finetuning.md), with rollouts isolated in [ACA Sandboxes](./backends/aca-sandboxes.md). | +| [`backends/`](./backends/) | Backend integrations: [SkyRL](./backends/skyrl.md) (open reference), [TRL](./backends/trl.md) (OpenEnv's recommended framework, with a runnable CPU [`trl_echo_demo.py`](./backends/trl_echo_demo.py)), [Tinker](./backends/tinker.md) (runnable remote LoRA [`tinker_echo_demo.py`](./backends/tinker_echo_demo.py)), [Foundry Fine-Tuning](./backends/foundry-finetuning.md), with rollouts isolated in [ACA Sandboxes](./backends/aca-sandboxes.md). | ## The full hybrid (GRPO + ECHO) and the real numbers This CPU demo isolates the **env-token loss** (the part that needs no GPU to be convincing). The full `L_GRPO(actions) + λ·L_env(obs)` loop — and the ~2.3×/pass@1-doubling results — needs a GPU; `echo_loss.py` is the exact -objective, and [`backends/`](./backends/) shows the one-line config on SkyRL / -Tinker / Foundry Fine-Tuning. To show ECHO end-to-end you set up a small model and fine-tune it -with this loss; start with [SkyRL](./backends/skyrl.md). +objective, and [`backends/`](./backends/) shows how it maps to SkyRL / Tinker / +Foundry Fine-Tuning. The runnable [Tinker example](./backends/tinker.md) sends the +verifier-free env-token objective to a remote LoRA trainer; use +[SkyRL](./backends/skyrl.md) for the open full-hybrid reference. diff --git a/examples/echo_world_model/backends/tinker.md b/examples/echo_world_model/backends/tinker.md index aa94ef9c8..2753d4744 100644 --- a/examples/echo_world_model/backends/tinker.md +++ b/examples/echo_world_model/backends/tinker.md @@ -1,42 +1,59 @@ # ECHO on Tinker -[Tinker](https://thinkingmachines.ai) (Thinking Machines) exposes four -primitives — `sample` / `forward_backward` / `optim_step` / `save_state` — and -lets *you* own the RL loop while the service owns the GPUs. ECHO fits its -`forward_backward` cleanly, because of one observation (also made by Prime -Intellect in *"True Agents Model the World"*): +[Tinker](https://tinker-docs.thinkingmachines.ai/tinker/) exposes the training +primitives while its service owns the GPUs. This runnable example sends the +existing terminal rollouts to a remote LoRA trainer and measures held-out +environment-token cross-entropy before and after training. -> **The env-token loss is just SFT on the observation tokens, and SFT is RL with -> a constant positive advantage.** So you don't need a second loss function — -> you reuse the exact same `forward_backward` and only change the **per-token -> advantage vector**. +## Run it -## The datum: one advantage per token +From `examples/echo_world_model`: -For each rollout you already build with [`trajectory.py`](../trajectory.py), -emit a per-token advantage: +```bash +export TINKER_API_KEY="..." +uv run backends/tinker_echo_demo.py --steps 15 +``` + +The inline dependency pins `tinker==0.22.7`. Tinker requires an account, API +key, and available credits. The default is a rank-16 LoRA on +`Qwen/Qwen3.5-4B`; use `--model` to select another supported public model. + +## How the role masks map to Tinker + +For each role-tagged rollout, the script tokenizes every segment and creates a +causal-language-model `Datum`: ```python -# action tokens -> GRPO group-relative advantage A_i (can be negative) -# env_output tokens (obs_mask) -> a constant positive advantage = λ (the world-model term) -# everything else (context, warnings) -> 0 -advantages = torch.zeros(T) -advantages[action_mask] = group_relative_advantage # standard GRPO -advantages[obs_mask] = world_model_coeff # ECHO, as constant +adv SFT +datum = tinker.types.Datum( + model_input=tinker.types.ModelInput.from_ints(token_ids[:-1]), + loss_fn_inputs={ + "target_tokens": token_ids[1:], + "weights": normalized_obs_mask[1:], + }, +) ``` -Then the usual Tinker step trains both at once: +The shift matters: input position `t` predicts target token `t + 1`, so the loss +weight comes from the **target** token's role. Only `env_output` has non-zero +weight; action, context, and warning tokens remain zero-weight conditioning +context. Batch-wide normalization makes Tinker's sum-reduced cross-entropy equal +the mean env-token CE used by the local demo. + +Each step follows Tinker's documented training path: ```python -datum = tinker.Datum(model_input=token_ids, loss_fn_inputs={"advantages": advantages, ...}) -training_client.forward_backward([datum], loss_fn="importance_sampling") -training_client.optim_step(adam_params) +fwdbwd = training_client.forward_backward(data, "cross_entropy") +optim = training_client.optim_step(tinker.types.AdamParams(learning_rate=2e-4)) +result = fwdbwd.result() +optim.result() ``` -Notes (from the Prime Intellect write-up): -- Skip KL / importance-ratio / icepop masking *on the SFT (obs) tokens* — they - are only needed for the RL (action) tokens. -- Normalize the RL and SFT token contributions **independently** so the dense - env tokens don't drown out the sparse action tokens. -- Keep λ small — they saw collapse at 0.05 for GLM-4.5-Air, stable at 0.005; - echo-rl's published Qwen3-8B config uses 0.05. Sweep it. +`TrainingClient.forward(..., "cross_entropy")` evaluates held-out data before +and after the loop without accumulating gradients. + +## Scope + +This is the existing **verifier-free** ECHO objective (`use_rl=False`), not the +full `L_GRPO + lambda * L_env` hybrid. It isolates the OpenEnv integration seam: +preserving token roles and putting cross-entropy weight only on environment +outputs. diff --git a/examples/echo_world_model/backends/tinker_echo_demo.py b/examples/echo_world_model/backends/tinker_echo_demo.py new file mode 100644 index 000000000..f7e57f92c --- /dev/null +++ b/examples/echo_world_model/backends/tinker_echo_demo.py @@ -0,0 +1,126 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "tinker==0.22.7", +# ] +# /// + +"""Train the verifier-free ECHO env-token objective with Tinker.""" + +from __future__ import annotations + +import argparse +import os +import sys +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +EXAMPLE_DIR = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(EXAMPLE_DIR)) + +from mini_terminal_env import TEST_TASKS, TRAIN_TASKS # noqa: E402 +from rollout import oracle_rollout # noqa: E402 +from trajectory import ENV_OUTPUT, Trajectory # noqa: E402 + +DEFAULT_MODEL = "Qwen/Qwen3.5-4B" +LORA_RANK = 16 +LEARNING_RATE = 2e-4 + + +def build_echo_data( + tinker_types: Any, + tokenizer: Any, + trajectories: Sequence[Trajectory], +) -> list[Any]: + """Tokenize rollouts and put loss weight only on shifted env-output tokens.""" + encoded: list[tuple[list[int], list[bool]]] = [] + env_token_count = 0 + for trajectory in trajectories: + token_ids: list[int] = [] + env_mask: list[bool] = [] + for segment in trajectory.segments: + segment_ids = tokenizer.encode(segment.text, add_special_tokens=False) + token_ids.extend(segment_ids) + env_mask.extend([segment.role == ENV_OUTPUT] * len(segment_ids)) + if len(token_ids) < 2: + raise ValueError("ECHO trajectories must contain at least two tokens") + shifted_env_mask = env_mask[1:] + env_token_count += sum(shifted_env_mask) + encoded.append((token_ids, shifted_env_mask)) + + if env_token_count == 0: + raise ValueError("ECHO data must contain at least one env_output target token") + + return [ + tinker_types.Datum( + model_input=tinker_types.ModelInput.from_ints(tokens=token_ids[:-1]), + loss_fn_inputs={ + "target_tokens": token_ids[1:], + "weights": [float(is_env) / env_token_count for is_env in env_mask], + }, + ) + for token_ids, env_mask in encoded + ] + + +def evaluate_env_ce(training_client: Any, data: Sequence[Any]) -> float: + """Return mean cross-entropy over the globally normalized env tokens.""" + result = training_client.forward(list(data), "cross_entropy").result() + weighted_logprob = 0.0 + for output, datum in zip(result.loss_fn_outputs, data, strict=True): + weighted_logprob += sum( + logprob * weight + for logprob, weight in zip( + output["logprobs"].tolist(), + datum.loss_fn_inputs["weights"].tolist(), + strict=True, + ) + ) + return -weighted_logprob + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--steps", type=int, default=15) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.steps < 1: + raise SystemExit("--steps must be at least 1") + if not os.getenv("TINKER_API_KEY"): + raise SystemExit("TINKER_API_KEY is required") + + import tinker + + print(f"Creating Tinker LoRA trainer for {args.model}") + training_client = tinker.ServiceClient().create_lora_training_client( + base_model=args.model, + rank=LORA_RANK, + ) + tokenizer = training_client.get_tokenizer() + train_data = build_echo_data( + tinker.types, tokenizer, [oracle_rollout(task) for task in TRAIN_TASKS] + ) + heldout_data = build_echo_data( + tinker.types, tokenizer, [oracle_rollout(task) for task in TEST_TASKS] + ) + + before = evaluate_env_ce(training_client, heldout_data) + for _ in range(args.steps): + backward = training_client.forward_backward(train_data, "cross_entropy") + optimizer = training_client.optim_step( + tinker.types.AdamParams(learning_rate=LEARNING_RATE) + ) + backward.result() + optimizer.result() + after = evaluate_env_ce(training_client, heldout_data) + + print(f"held-out env-token CE: {before:.3f} -> {after:.3f} nats/token") + + +if __name__ == "__main__": + main() diff --git a/examples/echo_world_model/test_tinker_echo.py b/examples/echo_world_model/test_tinker_echo.py new file mode 100644 index 000000000..d0ad87b07 --- /dev/null +++ b/examples/echo_world_model/test_tinker_echo.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +from backends.tinker_echo_demo import build_echo_data +from trajectory import ACTION, CONTEXT, ENV_OUTPUT, Segment, Trajectory + + +class _CharTokenizer: + def encode(self, text, add_special_tokens=False): + return [ord(char) for char in text] + + +@dataclass +class _ModelInput: + tokens: list[int] + + @classmethod + def from_ints(cls, tokens): + return cls(tokens) + + +@dataclass +class _Datum: + model_input: _ModelInput + loss_fn_inputs: dict + + +_TINKER_TYPES = SimpleNamespace(Datum=_Datum, ModelInput=_ModelInput) + + +def test_build_echo_data_shifts_and_normalizes_env_only_weights(): + trajectories = [ + Trajectory( + segments=[ + Segment(CONTEXT, "C"), + Segment(ACTION, "AB"), + Segment(ENV_OUTPUT, "XY"), + ], + reward=1.0, + ), + Trajectory([Segment(CONTEXT, "D"), Segment(ENV_OUTPUT, "Z")], reward=0.0), + ] + + first, second = build_echo_data(_TINKER_TYPES, _CharTokenizer(), trajectories) + + assert first.model_input.tokens == [ord(char) for char in "CABX"] + assert first.loss_fn_inputs["target_tokens"] == [ord(char) for char in "ABXY"] + assert first.loss_fn_inputs["weights"] == [0.0, 0.0, 1 / 3, 1 / 3] + assert second.loss_fn_inputs["weights"] == [1 / 3] + + +def test_build_echo_data_accepts_real_tinker_types(): + tinker = pytest.importorskip("tinker") + trajectory = Trajectory( + [Segment(CONTEXT, "prompt"), Segment(ENV_OUTPUT, "result")], reward=0.0 + ) + + [datum] = build_echo_data(tinker.types, _CharTokenizer(), [trajectory]) + + assert isinstance(datum, tinker.types.Datum) + assert datum.model_input.to_ints() == [ord(char) for char in "promptresul"] + assert sum(datum.loss_fn_inputs["weights"].tolist()) == pytest.approx(1.0) diff --git a/examples/echo_world_model/trajectory.py b/examples/echo_world_model/trajectory.py index bebfdbe4b..7ecdc9926 100644 --- a/examples/echo_world_model/trajectory.py +++ b/examples/echo_world_model/trajectory.py @@ -15,8 +15,10 @@ from __future__ import annotations from dataclasses import dataclass +from typing import TYPE_CHECKING -import torch +if TYPE_CHECKING: + import torch # token roles CONTEXT = "context" # system prompt / task — given, not a loss target @@ -59,6 +61,8 @@ def tokenize_trajectory( Returns dict of 1-D tensors: ``input_ids``, ``action_mask``, ``obs_mask``, ``warning_mask`` (the last three boolean, aligned to ``input_ids``). """ + import torch + ids: list[int] = [] roles: list[str] = [] for seg in traj.segments: