inductive-mlxrl 0.2.1__tar.gz → 0.2.2__tar.gz

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/CHANGELOG.md

@@ -5,6 +5,17 @@ All notable changes to this project are documented here.
 The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project uses semantic versioning once releases are cut.
 
+## [0.2.2] - 2026-06-25
+
+### Added
+
+- ECHO world-modeling support for tagged multi-turn trajectories, composing an
+  independently normalized ECHO SFT loss with the existing GiGPO action loss.
+- `echo_alpha` scheduling with constant and linear-taper-to-zero modes, plus
+  separate `loss_echo` and `echo_accuracy` diagnostics.
+- Real-Qwen smoke coverage for tagged ECHO batches on Qwen3-0.6B and
+  Qwen3.5-9B MLX 4-bit models.
+
 ## [0.2.1] - 2026-06-25
 
 ### Fixed

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inductive-mlxrl
-Version: 0.2.1
+Version: 0.2.2
 Summary: Small single-process RL post-training for LLMs on Apple Silicon with MLX.
 Project-URL: Homepage, https://github.com/inductiveML/mlxrl
 Project-URL: Repository, https://github.com/inductiveML/mlxrl
@@ -300,6 +300,27 @@ currently exact for token-mean policy losses: base GRPO, DAPO, GSPO token mode,
 RLOO, and Dr. GRPO with `loss_reduction="token_mean"`. Sequence-reduced losses
 should keep `micro_batch_size=0`.
 
+## ECHO World Modeling
+
+ECHO is an additive GiGPO trainer feature for tagged multi-turn trajectories.
+Producer code may attach `token_roles` to each `Trajectory`, aligned to
+`full_token_ids`: `0 = MASKED`, `1 = ACTION`, `2 = ECHO`. ACTION tokens keep
+the normal GiGPO policy-gradient loss. ECHO tokens add an independently
+normalized SFT term, `alpha * NLL`, in the same forward/backward pass.
+
+`echo_alpha` defaults to `0.0`, which leaves current GiGPO behavior untouched.
+For early world-modeling, start small, roughly `0.005` to `0.05`, and consider
+`echo_schedule = "linear_taper_to_zero"` with `echo_taper_steps = N` so training
+returns to pure RL. `alpha` is model-size-sensitive; too-high values can
+overfit tool responses or collapse RL.
+
+ECHO is content-agnostic by design. `mlxrl` never inspects token text to decide
+what is an environment or tool response. The rollout producer owns the tags:
+MASKED tokens have zero gradient, ACTION tokens receive the RL advantage, and
+ECHO tokens receive only the SFT loss with no KL penalty or importance ratio.
+The trainer logs `loss_echo` and token-level `echo_acc` separately from the
+ACTION loss.
+
 ## Policy Semantics
 
 - The base model is frozen before LoRA injection.

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/README.md

@@ -275,6 +275,27 @@ currently exact for token-mean policy losses: base GRPO, DAPO, GSPO token mode,
 RLOO, and Dr. GRPO with `loss_reduction="token_mean"`. Sequence-reduced losses
 should keep `micro_batch_size=0`.
 
+## ECHO World Modeling
+
+ECHO is an additive GiGPO trainer feature for tagged multi-turn trajectories.
+Producer code may attach `token_roles` to each `Trajectory`, aligned to
+`full_token_ids`: `0 = MASKED`, `1 = ACTION`, `2 = ECHO`. ACTION tokens keep
+the normal GiGPO policy-gradient loss. ECHO tokens add an independently
+normalized SFT term, `alpha * NLL`, in the same forward/backward pass.
+
+`echo_alpha` defaults to `0.0`, which leaves current GiGPO behavior untouched.
+For early world-modeling, start small, roughly `0.005` to `0.05`, and consider
+`echo_schedule = "linear_taper_to_zero"` with `echo_taper_steps = N` so training
+returns to pure RL. `alpha` is model-size-sensitive; too-high values can
+overfit tool responses or collapse RL.
+
+ECHO is content-agnostic by design. `mlxrl` never inspects token text to decide
+what is an environment or tool response. The rollout producer owns the tags:
+MASKED tokens have zero gradient, ACTION tokens receive the RL advantage, and
+ECHO tokens receive only the SFT loss with no KL penalty or importance ratio.
+The trainer logs `loss_echo` and token-level `echo_acc` separately from the
+ACTION loss.
+
 ## Policy Semantics
 
 - The base model is frozen before LoRA injection.

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/__init__.py

@@ -2,4 +2,4 @@
 
 __all__ = ["__version__"]
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/algorithm.py

@@ -23,6 +23,9 @@ class AlgorithmLossMetrics:
     kl: mx.array
     mean_ratio: mx.array
     clip_fraction: mx.array
+    loss_action: mx.array | None = None
+    loss_echo: mx.array | None = None
+    echo_accuracy: mx.array | None = None
 
 
 class Algorithm(Protocol):

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/cli.py

@@ -38,6 +38,7 @@ from mlxrl.data.gsm8k import (
     format_gsm8k_prompt,
 )
 from mlxrl.data.rewards import accuracy_reward, format_reward
+from mlxrl.echo import EchoSchedule, EchoScheduleName
 from mlxrl.env import EnvFactory, RecurringStateTextEnv, SingleTurnRewardEnv
 from mlxrl.policy.logprobs import pad_token_id_from_tokenizer
 from mlxrl.policy.model import (
@@ -481,6 +482,9 @@ def _namespace_from_train_config(config: TrainConfig) -> argparse.Namespace:
         max_turns=config.max_turns,
         rollout_mode=config.rollout_mode,
         env_name=config.env_name,
+        echo_alpha=config.echo_alpha,
+        echo_schedule=config.echo_schedule,
+        echo_taper_steps=config.echo_taper_steps,
     )
 
 
@@ -500,6 +504,12 @@ def _apply_train_overrides(config: TrainConfig, args: argparse.Namespace) -> Tra
         replacements["seed"] = args.seed
     if args.output is not None:
         replacements["output"] = args.output
+    if args.echo_alpha is not None:
+        replacements["echo_alpha"] = args.echo_alpha
+    if args.echo_schedule is not None:
+        replacements["echo_schedule"] = args.echo_schedule
+    if args.echo_taper_steps is not None:
+        replacements["echo_taper_steps"] = args.echo_taper_steps
     if args.algorithm is not None:
         replacements["algorithm"] = replace(
             config.algorithm,
@@ -575,6 +585,11 @@ def _phase_agentic_train(config: TrainConfig) -> int:
     )
     algorithm = _trajectory_algorithm_from_config(config)
     env_factory = _env_factory_from_config(config)
+    echo_schedule = EchoSchedule(
+        alpha=config.echo_alpha,
+        schedule=cast(EchoScheduleName, config.echo_schedule),
+        taper_steps=config.echo_taper_steps,
+    )
     reward_history: list[float] = []
     kl_history: list[float] = []
     final_batch = None
@@ -604,6 +619,7 @@ def _phase_agentic_train(config: TrainConfig) -> int:
             use_checkpoint=config.gradient_checkpointing,
             compute_reference=config.kl_beta != 0.0,
         )
+        echo_alpha = echo_schedule.value(step)
         metrics = optimizer_step_trajectory(
             model=model,
             optimizer=optimizer,
@@ -613,6 +629,7 @@ def _phase_agentic_train(config: TrainConfig) -> int:
             algorithm=algorithm,
             use_checkpoint=config.gradient_checkpointing,
             micro_batch_size=config.micro_batch_size,
+            echo_alpha=echo_alpha,
         )
         reward_history.append(metrics.mean_reward)
         kl_history.append(metrics.kl)
@@ -624,6 +641,9 @@ def _phase_agentic_train(config: TrainConfig) -> int:
             f"kl={metrics.kl:.6f} "
             f"ratio={metrics.mean_ratio:.6f} "
             f"clip={metrics.clip_fraction:.6f} "
+            f"echo_alpha={metrics.echo_alpha:.6f} "
+            f"loss_echo={metrics.loss_echo:.6f} "
+            f"echo_acc={metrics.echo_accuracy:.6f} "
             f"loss={metrics.loss:.6f}",
             flush=True,
         )
@@ -645,7 +665,14 @@ def _phase_agentic_train(config: TrainConfig) -> int:
         old_policy_logprobs=final_batch.old_policy_logprobs,
         reference_logprobs=final_batch.reference_logprobs,
         action_mask=final_batch.action_mask,
+        echo_mask=(
+            final_batch.echo_mask
+            if final_batch.echo_mask is not None
+            else mx.zeros_like(final_batch.action_mask)
+        ),
         loss=mx.array(final_metrics.loss, dtype=mx.float32),
+        loss_echo=mx.array(final_metrics.loss_echo, dtype=mx.float32),
+        echo_accuracy=mx.array(final_metrics.echo_accuracy, dtype=mx.float32),
         kl=mx.array(final_metrics.kl, dtype=mx.float32),
     )
     print(f"reference_output: {output}")
@@ -881,6 +908,13 @@ def build_parser() -> argparse.ArgumentParser:
     train.add_argument("--available-memory-gb", type=float, default=None)
     train.add_argument("--auto-fit", action="store_true")
     train.add_argument("--dry-run", action="store_true")
+    train.add_argument("--echo-alpha", type=float, default=None)
+    train.add_argument(
+        "--echo-schedule",
+        choices=["constant", "linear_taper_to_zero"],
+        default=None,
+    )
+    train.add_argument("--echo-taper-steps", type=int, default=None)
     train.set_defaults(func=_phase_train)
 
     phase2 = subparsers.add_parser(

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/config.py

@@ -7,7 +7,9 @@ import re
 import tomllib
 from dataclasses import dataclass, field, replace
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
+
+from mlxrl.echo import EchoSchedule, EchoScheduleName
 
 
 class ConfigError(ValueError):
@@ -107,6 +109,9 @@ class TrainConfig:
     iogpu_wired_limit_mb: int | None = None
     micro_batch_size: int = 0
     use_chat_template: bool = False
+    echo_alpha: float = 0.0
+    echo_schedule: str = "constant"
+    echo_taper_steps: int | None = None
     output: str = "reference_outputs/phase1_reference.npz"
     optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
     sampling: SamplingConfigData = field(default_factory=SamplingConfigData)
@@ -147,6 +152,16 @@ class TrainConfig:
             raise ConfigError("iogpu_wired_limit_mb must be positive or null.")
         if self.micro_batch_size < 0:
             raise ConfigError("micro_batch_size must be non-negative.")
+        if self.echo_alpha < 0:
+            raise ConfigError("echo_alpha must be non-negative.")
+        try:
+            EchoSchedule(
+                alpha=self.echo_alpha,
+                schedule=cast(EchoScheduleName, self.echo_schedule),
+                taper_steps=self.echo_taper_steps,
+            )
+        except ValueError as error:
+            raise ConfigError(str(error)) from error
         validate_output_path(self.output)
         self.optimizer.validate()
         self.sampling.validate()

inductive_mlxrl-0.2.2/mlxrl/echo.py

@@ -0,0 +1,43 @@
+"""ECHO token-role helpers and alpha scheduling."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+MASKED = 0
+ACTION = 1
+ECHO = 2
+VALID_TOKEN_ROLES = frozenset({MASKED, ACTION, ECHO})
+EchoScheduleName = Literal["constant", "linear_taper_to_zero"]
+
+
+@dataclass(frozen=True)
+class EchoSchedule:
+    """Schedule ECHO's SFT strength without changing the RL objective."""
+
+    alpha: float = 0.0
+    schedule: EchoScheduleName = "constant"
+    taper_steps: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.alpha < 0:
+            raise ValueError("ECHO alpha must be non-negative.")
+        if self.schedule not in {"constant", "linear_taper_to_zero"}:
+            raise ValueError("ECHO schedule must be 'constant' or 'linear_taper_to_zero'.")
+        if (
+            self.schedule == "linear_taper_to_zero"
+            and (self.taper_steps is None or self.taper_steps < 1)
+        ):
+            raise ValueError("linear_taper_to_zero requires taper_steps >= 1.")
+
+    def value(self, step: int) -> float:
+        """Return alpha for a zero-based optimizer step."""
+
+        if step < 0:
+            raise ValueError("step must be non-negative.")
+        if self.schedule == "constant" or self.alpha == 0.0:
+            return self.alpha
+        assert self.taper_steps is not None
+        fraction = max(0.0, 1.0 - (float(step) / float(self.taper_steps)))
+        return self.alpha * fraction

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/policy/__init__.py

@@ -21,10 +21,14 @@ from mlxrl.policy.model import (
     strict_lora_config,
 )
 from mlxrl.policy.trajectory_logprobs import (
+    TaggedTrajectoryDualLogprobs,
+    TaggedTrajectoryLogprobs,
     TrajectoryDualLogprobs,
     TrajectoryLogprobs,
     trajectory_action_logprobs,
     trajectory_dual_logprobs,
+    trajectory_tagged_dual_logprobs,
+    trajectory_tagged_logprobs,
 )
 
 __all__ = [
@@ -34,6 +38,8 @@ __all__ = [
     "DualLogprobs",
     "LoRAConfig",
     "Phase0Report",
+    "TaggedTrajectoryDualLogprobs",
+    "TaggedTrajectoryLogprobs",
     "TrajectoryDualLogprobs",
     "TrajectoryLogprobs",
     "adapters_disabled",
@@ -48,4 +54,6 @@ __all__ = [
     "strict_lora_config",
     "trajectory_action_logprobs",
     "trajectory_dual_logprobs",
+    "trajectory_tagged_dual_logprobs",
+    "trajectory_tagged_logprobs",
 ]

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/policy/trajectory_logprobs.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass
 import mlx.core as mx
 import mlx.nn as nn
 
+from mlxrl.echo import ACTION, ECHO
 from mlxrl.policy.logprobs import adapters_disabled, target_logprobs_from_logits
 from mlxrl.trajectory import Trajectory
 
@@ -29,6 +30,28 @@ class TrajectoryDualLogprobs:
     mask: mx.array
 
 
+@dataclass(frozen=True)
+class TaggedTrajectoryLogprobs:
+    """Full-forward logprobs and masks for tagged trajectory target tokens."""
+
+    logprobs: mx.array
+    target_ids: mx.array
+    predictions: mx.array
+    action_mask: mx.array
+    echo_mask: mx.array
+
+
+@dataclass(frozen=True)
+class TaggedTrajectoryDualLogprobs:
+    """Policy/reference logprobs for tagged trajectory target tokens."""
+
+    policy: mx.array
+    reference: mx.array
+    target_ids: mx.array
+    action_mask: mx.array
+    echo_mask: mx.array
+
+
 def trajectory_action_logprobs(
     model: nn.Module,
     trajectories: Sequence[Trajectory],
@@ -83,6 +106,100 @@ def trajectory_action_logprobs(
     )
 
 
+def trajectory_tagged_logprobs(
+    model: nn.Module,
+    trajectories: Sequence[Trajectory],
+    pad_token_id: int = 0,
+    use_checkpoint: bool = False,
+) -> TaggedTrajectoryLogprobs:
+    """Gather logprobs for ACTION and ECHO tags from one full forward."""
+
+    del use_checkpoint
+    if not trajectories:
+        raise ValueError("At least one trajectory is required.")
+    max_sequence_len = max(len(trajectory.full_token_ids) for trajectory in trajectories)
+    if max_sequence_len < 2:
+        raise ValueError("A trajectory needs at least two tokens for logprobs.")
+
+    target_width = max_sequence_len - 1
+    input_rows: list[list[int]] = []
+    target_rows: list[list[int]] = []
+    action_mask_rows: list[list[float]] = []
+    echo_mask_rows: list[list[float]] = []
+    for trajectory in trajectories:
+        sequence = list(trajectory.full_token_ids)
+        roles = trajectory.token_roles_or_default()
+        target_roles = list(roles[1:])
+        pad_count = target_width - (len(sequence) - 1)
+        input_rows.append(sequence[:-1] + [pad_token_id] * pad_count)
+        target_rows.append(sequence[1:] + [pad_token_id] * pad_count)
+        action_mask_rows.append(
+            [1.0 if role == ACTION else 0.0 for role in target_roles]
+            + [0.0] * pad_count
+        )
+        echo_mask_rows.append(
+            [1.0 if role == ECHO else 0.0 for role in target_roles]
+            + [0.0] * pad_count
+        )
+
+    input_ids = mx.array(input_rows, dtype=mx.int32)
+    target_ids = mx.array(target_rows, dtype=mx.int32)
+    logits = model(input_ids)
+    logprobs = target_logprobs_from_logits(logits, target_ids)
+    predictions = mx.argmax(logits, axis=-1).astype(mx.int32)
+    return TaggedTrajectoryLogprobs(
+        logprobs=logprobs,
+        target_ids=target_ids,
+        predictions=predictions,
+        action_mask=mx.array(action_mask_rows, dtype=mx.float32),
+        echo_mask=mx.array(echo_mask_rows, dtype=mx.float32),
+    )
+
+
+def trajectory_tagged_dual_logprobs(
+    model: nn.Module,
+    trajectories: Sequence[Trajectory],
+    pad_token_id: int = 0,
+    use_checkpoint: bool = False,
+    compute_reference: bool = True,
+) -> TaggedTrajectoryDualLogprobs:
+    """Compute tagged policy/reference logprobs using one model object."""
+
+    policy = trajectory_tagged_logprobs(
+        model,
+        trajectories,
+        pad_token_id=pad_token_id,
+        use_checkpoint=use_checkpoint,
+    )
+    if not compute_reference:
+        return TaggedTrajectoryDualLogprobs(
+            policy=policy.logprobs,
+            reference=mx.stop_gradient(policy.logprobs),
+            target_ids=policy.target_ids,
+            action_mask=policy.action_mask,
+            echo_mask=policy.echo_mask,
+        )
+    with adapters_disabled(model):
+        reference = trajectory_tagged_logprobs(
+            model,
+            trajectories,
+            pad_token_id=pad_token_id,
+            use_checkpoint=use_checkpoint,
+        )
+        mx.eval(  # Reference sync: materialize tagged ref before restoring adapters.
+            reference.logprobs,
+            reference.action_mask,
+            reference.echo_mask,
+        )
+    return TaggedTrajectoryDualLogprobs(
+        policy=policy.logprobs,
+        reference=reference.logprobs,
+        target_ids=policy.target_ids,
+        action_mask=policy.action_mask,
+        echo_mask=policy.echo_mask,
+    )
+
+
 def trajectory_dual_logprobs(
     model: nn.Module,
     trajectories: Sequence[Trajectory],

{inductive_mlxrl-0.2.1 → inductive_mlxrl-0.2.2}/mlxrl/train/grpo.py

@@ -50,6 +50,10 @@ class StepMetrics:
     mean_ratio: float
     clip_fraction: float
     mean_reward: float
+    loss_action: float = 0.0
+    loss_echo: float = 0.0
+    echo_accuracy: float = 0.0
+    echo_alpha: float = 0.0
 
 
 def _slice_batch(batch: GRPOBatch, start: int, end: int) -> GRPOBatch: