Add Tinker ECHO example

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.

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.

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()

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)

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: