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

inductive_mlxrl-0.2.2/CHANGELOG.md

@@ -0,0 +1,87 @@
+# Changelog
+
+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
+
+- Keep model-backed agentic/GiGPO decode tokens rank-2 during KV-cache decode,
+  fixing real Qwen3 model rollouts.
+- Use valid repo-relative output paths in the committed Qwen smoke configs.
+- Preserve eval-mode old-policy/reference logprob prep while restoring caller
+  train mode afterward.
+
+## [0.2.0] - 2026-06-24
+
+### Added
+
+- Multi-turn agentic rollout path with Gym-style `Environment` / `EnvFactory`
+  protocols and deterministic reference environments.
+- Trajectory data model with action-token spans and full-forward trajectory
+  logprob gathering for 4-bit policy, old-policy, and reference semantics.
+- GiGPO (`Group-in-Group Policy Optimization`) with episode-level and
+  anchor-state step-level advantages, plus an exact `omega=0` reduction gate.
+- Trajectory batch builder and optimizer step that preserve the existing
+  single-turn GRPO-family APIs.
+- Config-driven GiGPO training route with `max_turns`, `rollout_mode`,
+  `env_name`, `max_observation_len`, `gigpo_omega`, `gigpo_gamma`, and
+  normalization knobs.
+
+### Changed
+
+- Extend the memory estimator to label multi-turn configs as estimated
+  OOM-risk shapes while preserving the measured v0.1 single-turn anchors.
+
+## [0.1.2] - 2026-06-24
+
+### Fixed
+
+- Move the Python `<3.15` ceiling from package-level `Requires-Python` to the
+  MLX runtime dependency markers so universal resolvers can add the package to
+  projects with broad Python ranges while skipping unsupported MLX splits.
+
+## [0.1.1] - 2026-06-24
+
+### Fixed
+
+- Mark `mlx` and `mlx-lm` runtime dependencies as macOS arm64-only and bound the
+  package to Python 3.11 through 3.14 so universal resolvers do not try to solve
+  unsupported MLX platform or Python splits.
+
+## [0.1.0] - 2026-06-24
+
+### Added
+
+- Single-process MLX GRPO+QLoRA training path for Apple Silicon.
+- Prefix-cached grouped rollout engine with compiled decode support.
+- Adapter-disabled reference logprob path on one model object.
+- Full-forward old-policy logprob recompute for 4-bit correctness.
+- Algorithm protocol with GRPO, Dr. GRPO, DAPO, GSPO, and RLOO.
+- DAPO `filter_batch` hook for dynamic zero-advantage group filtering.
+- Per-layer MLX-LM gradient checkpointing for DeltaNet/linear-attention models.
+- Micro-batched gradient accumulation for token-mean losses.
+- Typed config schema and memory preflight helper calibrated to measured anchors.
+- Phase 4 benchmark harness for `mlxrl`, `mlx-lm`, `mlx-tune`, and `mlx-lm-lora`.
+- Linux CPU-safe CI for docs/config/import-direction checks.
+
+### Notes
+
+- `mlxrl` is pre-1.0; public APIs may change while the correctness gates settle.
+- Full MLX/Metal tests require Apple Silicon or a self-hosted Mac runner.
+- Published on PyPI as `inductive-mlxrl`; the import package and CLI command
+  remain `mlxrl`.

{inductive_mlxrl-0.1.1 → inductive_mlxrl-0.2.2}/CONTRIBUTING.md

@@ -20,7 +20,9 @@ estimator anchor tests are part of the contract.
 
 ## Adding An Algorithm
 
-Use `RLOOAlgorithm` in `mlxrl/algo/grpo.py` as the smallest template.
+Use `RLOOAlgorithm` in `mlxrl/algo/grpo.py` as the smallest single-turn
+template. Use `GiGPOAlgorithm` in `mlxrl/algo/gigpo.py` as the trajectory
+template.
 
 Required steps:
 
@@ -35,6 +37,23 @@ Required steps:
 If an algorithm needs to drop or reshape examples before the loss, use the
 `filter_batch` hook. Do not special-case it in the trainer.
 
+Trajectory algorithms implement `TrajectoryAlgorithm`, ship a hand-computed
+multi-turn loss/gradient test, and prove their reduction to an existing
+objective when a degenerate config should do so. They may consume trajectory
+metadata such as env state ids, but `rollout/`, `policy/`, and `train/` still
+must not import concrete algorithms.
+
+## Adding An Environment
+
+External envs integrate through `EnvFactory(task, seed, group_index) ->
+Environment`. Implement `reset()`, `step(action)`, `max_turns`, and
+`state_id(observation)`. The env owns the anchor-state definition; the rollout
+engine only records ids, and GiGPO groups them later.
+
+Keep envs deterministic under the provided seed when possible. Add a tiny test
+that resets, steps, returns stable state ids, and exercises truncation or
+termination. Do not add omega gym itself to this repo.
+
 ## MLX And CI
 
 GitHub-hosted Linux runners cannot run the MLX/Metal path. Linux CI gates the
@@ -47,9 +66,11 @@ would close this coverage gap.
 
 ## Package Compatibility
 
-The PyPI distribution supports Python 3.11 through 3.14. Keep the `mlx` and
-`mlx-lm` runtime dependencies marked for macOS arm64 so universal resolvers do
-not try to solve MLX wheels for unsupported platforms.
+The PyPI distribution requires Python 3.11+. The MLX runtime path is supported
+on Apple Silicon with Python 3.11 through 3.14. Keep the `mlx` and `mlx-lm`
+runtime dependencies marked for macOS arm64 and Python `<3.15` so universal
+resolvers do not try to solve MLX wheels for unsupported platform or Python
+splits.
 
 ## Release Cycle
 
@@ -103,6 +124,7 @@ In scope:
 - single-process Apple Silicon RL post-training;
 - QLoRA on local MLX LLMs;
 - critic-free on-policy algorithms;
+- multi-turn agentic environments through the small env protocol;
 - memory-conscious rollout and training paths.
 
 Out of scope:

{inductive_mlxrl-0.1.1 → inductive_mlxrl-0.2.2}/DESIGN.md

@@ -7,10 +7,10 @@ fast enough that rollout is the main problem and the loss stays thin.
 ## Critic-Free By Design
 
 The library is intentionally limited to rollout-based policy-gradient methods:
-GRPO, Dr. GRPO, DAPO, GSPO, and RLOO. PPO is out of scope because it needs a
-critic/value model path, a different forward pass, and a different memory
-profile. DPO and ORPO are also out of scope because they are offline preference
-objectives rather than on-policy rollout algorithms.
+GRPO, Dr. GRPO, DAPO, GSPO, RLOO, and GiGPO. PPO is out of scope because it
+needs a critic/value model path, a different forward pass, and a different
+memory profile. DPO and ORPO are also out of scope because they are offline
+preference objectives rather than on-policy rollout algorithms.
 
 That boundary matters. `mlxrl` keeps one policy model object in memory, attaches
 QLoRA adapters to it, and computes reference logprobs by disabling those
@@ -32,6 +32,24 @@ stable while algorithms change. DAPO's dynamic sampling is the proof that the
 interface is more general than "GRPO with renamed constants": it drops
 zero-advantage groups through `filter_batch` without special trainer branches.
 
+Multi-turn algorithms use the sibling `TrajectoryAlgorithm` protocol rather
+than widening the single-turn `Algorithm` contract. GiGPO is the first
+implementation: the rollout engine records env-provided state ids, while the
+algorithm owns anchor-state grouping and credit assignment.
+
+## Agentic Environment Seam
+
+Agentic rollout is built around `EnvFactory(task, seed, group_index) ->
+Environment`. An environment exposes `reset()`, `step(action)`, `max_turns`,
+and `state_id(observation)`. The state id is deliberately supplied by the env:
+omega gym, WebShop, ALFWorld, or a toy env each decides what counts as the same
+anchor state.
+
+The engine records observations, actions, rewards, state ids, full token
+history, and action-token spans. It does not compute anchor groups. This keeps
+rollout reusable for future trajectory algorithms while GiGPO performs the
+group-in-group advantage calculation offline from the same trajectories.
+
 ## Correctness Gates
 
 Speed changes are allowed only behind equivalence gates. The original Phase 1
@@ -40,6 +58,11 @@ token-for-token at fixed seed and match loss within tolerance. The protocol
 refactor was checked against the pre-refactor GRPO and Dr. GRPO losses with
 zero loss and adapter-gradient difference.
 
+GiGPO adds two reduction gates: `omega=0` must match episode-level GRPO over the
+same trajectories, and single-turn one-step trajectories with `omega=0` must
+match the v0.1 GRPO path through loss and adapter gradient. Toy multi-turn
+cases also carry hand-computed advantage, loss, and gradient checks.
+
 Old-policy changes use a stronger gate than naive equality. At 4-bit, the
 rollout-time cached decode realization and the full-forward realization are not
 numerically identical. Stored rollout logprobs are attractive because they are
@@ -56,6 +79,11 @@ for gradient-bearing or importance-weighting quantities. Anything used in the
 loss denominator, KL, or adapter gradient must come from the full-forward path
 unless a dedicated stability gate proves otherwise.
 
+For multi-turn trajectories, this means policy, old-policy, and reference
+logprobs are gathered from a full forward over the complete token history, then
+masked down to action-token spans only. Rollout-time cached logprobs remain
+diagnostic data, not training data.
+
 ## Memory As A First-Class Constraint
 
 Apple Silicon's unified memory is the deployment target, not an afterthought.
@@ -70,3 +98,9 @@ anchors and labels long-sequence uncheckpointed hybrid configs as OOM-risk
 estimates, not measurements. Its purpose is to nudge users toward the knob that
 will most likely make a run fit: enable checkpointing, reduce G, then reduce
 completion length.
+
+Multi-turn configs add trajectory length as `prompt + turns * action_tokens +
+(turns - 1) * observation_tokens`. Parallel-per-turn rollout keeps G
+per-trajectory cache/state objects resident; sequential mode is the lower-memory
+fallback. Until more measured anchors exist, multi-turn estimates are labeled as
+OOM-risk estimates rather than confident measurements.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inductive-mlxrl
-Version: 0.1.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
@@ -18,23 +18,25 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: <3.15,>=3.11
-Requires-Dist: mlx-lm>=0.31.3; platform_system == 'Darwin' and platform_machine == 'arm64'
-Requires-Dist: mlx>=0.31.2; platform_system == 'Darwin' and platform_machine == 'arm64'
+Requires-Python: >=3.11
+Requires-Dist: mlx-lm>=0.31.3; python_version < '3.15' and platform_system == 'Darwin' and platform_machine == 'arm64'
+Requires-Dist: mlx>=0.31.2; python_version < '3.15' and platform_system == 'Darwin' and platform_machine == 'arm64'
 Description-Content-Type: text/markdown
 
 # mlxrl
 
-Fast on-policy MLX RL for Apple Silicon; not a general RL framework, not
-preference tuning, and not distributed training.
+Fast on-policy MLX RL for Apple Silicon, including single-turn GRPO-family RL
+and multi-turn GiGPO; not a general RL framework, not preference tuning, and
+not distributed training.
 
 `mlxrl` is a small, single-process RL post-training library for LLMs on Apple
 Silicon. It is built around one idea: GRPO on MLX should be a fast batched
 rollout path with a thin loss and optimizer step on top, not a framework.
 
-The current implementation targets QLoRA GRPO on local 4-bit MLX models. It
-reuses `mlx-lm` model loading, LoRA layers, KV caches, and sampling utilities,
-and keeps generation and training in one Python process with one model object.
+The current implementation targets QLoRA GRPO/GiGPO on local 4-bit MLX models.
+It reuses `mlx-lm` model loading, LoRA layers, KV caches, and sampling
+utilities, and keeps generation and training in one Python process with one
+model object.
 
 `mlxrl` is pre-1.0. The correctness gates are stable, but import APIs and config
 fields may change before a 1.0 release.
@@ -66,6 +68,9 @@ UV_CACHE_DIR=.uv-cache uv run mlxrl train \
 - Adapter-disabled reference policy on the same model object.
 - GRPO, Dr. GRPO, DAPO, and GSPO loss variants.
 - RLOO (REINFORCE Leave-One-Out) as a critic-free rollout objective.
+- Multi-turn agentic rollout through an `Environment` / `EnvFactory` protocol.
+- GiGPO with episode-level and anchor-state step-level advantages.
+- Full-forward trajectory action-logprob recompute for 4-bit correctness.
 - QLoRA injection on dense and heterogeneous/hybrid layer stacks.
 - Qwen3.5-style hybrid support via MLX-LM auto LoRA targeting, including
   DeltaNet `linear_attn.in_proj_*` and dense attention `q/k/v/o_proj`.
@@ -101,10 +106,11 @@ Run commands through the local environment:
 UV_CACHE_DIR=.uv-cache uv run mlxrl --help
 ```
 
-Python 3.11 through 3.14 is supported. Runtime dependencies are intentionally
-small: `mlx` and `mlx-lm` on macOS arm64, matching the MLX wheel support window.
-Development dependencies include `pytest`, `ruff`, `pyright`, `mlx-tune`, and
-`mlx-lm-lora` for comparison benchmarks.
+Python 3.11+ is required. The MLX runtime path is supported on Apple Silicon
+with Python 3.11 through 3.14, matching the current MLX wheel support window.
+Runtime dependencies are intentionally small: `mlx` and `mlx-lm` on macOS
+arm64. Development dependencies include `pytest`, `ruff`, `pyright`,
+`mlx-tune`, and `mlx-lm-lora` for comparison benchmarks.
 The PyPI distribution name is `inductive-mlxrl`; the import package and console
 script remain `mlxrl`.
 
@@ -158,6 +164,14 @@ UV_CACHE_DIR=.uv-cache uv run mlxrl train \
   --available-memory-gb 48
 ```
 
+Reference multi-turn GiGPO run against the built-in recurring-state env:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl train \
+  --config examples/qwen3_0_6b_gigpo_recurring.toml \
+  --available-memory-gb 48
+```
+
 The config schema validates model id, quant bits, group size, completion/prompt
 lengths, checkpointing granularity, `iogpu.wired_limit_mb`, optimizer settings,
 algorithm hyperparameters, KL beta, and seed before a model is loaded. CLI
@@ -242,11 +256,71 @@ metrics = optimizer_step(
 )
 ```
 
+Multi-turn GiGPO uses the additive trajectory path:
+
+```python
+from mlxrl.algo import GiGPOAlgorithm
+from mlxrl.env import RecurringStateTextEnv
+from mlxrl.rollout import generate_agentic_trajectories
+from mlxrl.train import batch_from_trajectories, optimizer_step_trajectory
+
+algorithm = GiGPOAlgorithm(omega=1.0, gamma=1.0)
+trajectories = generate_agentic_trajectories(
+    model=model,
+    tokenizer=tokenizer,
+    env_factory=lambda task, seed, group: RecurringStateTextEnv(
+        task=str(task),
+        max_turns=3,
+    ),
+    tasks=["finish the task"],
+    group_size=4,
+    sampling=sampling,
+    rollout_mode="parallel_per_turn",
+)
+batch = batch_from_trajectories(
+    model=model,
+    trajectories=trajectories,
+    group_size=4,
+    pad_token_id=pad_token_id,
+    algorithm=algorithm,
+    compute_reference=beta != 0.0,
+)
+metrics = optimizer_step_trajectory(
+    model=model,
+    optimizer=optimizer,
+    batch=batch,
+    beta=beta,
+    pad_token_id=pad_token_id,
+    algorithm=algorithm,
+)
+```
+
 `micro_batch_size=0` keeps the original whole-batch path. Micro-batching is
 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.
@@ -277,6 +351,7 @@ import concrete algorithm implementations.
 | DAPO | asymmetric low/high clipping plus optional dynamic zero-advantage group filtering |
 | GSPO | sequence-level, length-normalized importance ratio and clipping |
 | RLOO | leave-one-out group baseline, no critic, no std-normalized advantage |
+| GiGPO | multi-turn episode advantage plus anchor-state step advantage |
 
 ## Memory Preflight
 
@@ -295,6 +370,9 @@ Qwen3.5-9B/G2/seq609/per-layer-checkpointed, `45.9 GB` for
 Qwen3.5-9B/G4/seq609/per-layer-checkpointed, and `36 GB` for
 Qwen3.5-9B/G2/seq128/no-checkpoint. For hybrid 9B no-checkpoint long-sequence
 configs, it reports an OOM-risk lower bound rather than a fake precise peak.
+Multi-turn configs include `max_turns` and `max_observation_len` in the
+estimated trajectory length and are labeled as estimated OOM-risk shapes until
+new measured anchors are added.
 For an obviously too-large Qwen3.5-9B/G8/prompt97/T512/no-checkpoint config on
 48 GB, it flags the run and suggests the measured-boundary fallback around
 G4/T512/checkpointed.
@@ -374,6 +452,7 @@ updates, or benchmark timing boundaries.
 ```text
 mlxrl/
   rollout/   # batched group generation
+  env/       # agentic environment protocols and reference envs
   policy/    # model loading, LoRA setup, logprob passes
   algo/      # GRPO-family advantages and losses
   train/     # value_and_grad and optimizer integration

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

@@ -1,15 +1,17 @@
 # mlxrl
 
-Fast on-policy MLX RL for Apple Silicon; not a general RL framework, not
-preference tuning, and not distributed training.
+Fast on-policy MLX RL for Apple Silicon, including single-turn GRPO-family RL
+and multi-turn GiGPO; not a general RL framework, not preference tuning, and
+not distributed training.
 
 `mlxrl` is a small, single-process RL post-training library for LLMs on Apple
 Silicon. It is built around one idea: GRPO on MLX should be a fast batched
 rollout path with a thin loss and optimizer step on top, not a framework.
 
-The current implementation targets QLoRA GRPO on local 4-bit MLX models. It
-reuses `mlx-lm` model loading, LoRA layers, KV caches, and sampling utilities,
-and keeps generation and training in one Python process with one model object.
+The current implementation targets QLoRA GRPO/GiGPO on local 4-bit MLX models.
+It reuses `mlx-lm` model loading, LoRA layers, KV caches, and sampling
+utilities, and keeps generation and training in one Python process with one
+model object.
 
 `mlxrl` is pre-1.0. The correctness gates are stable, but import APIs and config
 fields may change before a 1.0 release.
@@ -41,6 +43,9 @@ UV_CACHE_DIR=.uv-cache uv run mlxrl train \
 - Adapter-disabled reference policy on the same model object.
 - GRPO, Dr. GRPO, DAPO, and GSPO loss variants.
 - RLOO (REINFORCE Leave-One-Out) as a critic-free rollout objective.
+- Multi-turn agentic rollout through an `Environment` / `EnvFactory` protocol.
+- GiGPO with episode-level and anchor-state step-level advantages.
+- Full-forward trajectory action-logprob recompute for 4-bit correctness.
 - QLoRA injection on dense and heterogeneous/hybrid layer stacks.
 - Qwen3.5-style hybrid support via MLX-LM auto LoRA targeting, including
   DeltaNet `linear_attn.in_proj_*` and dense attention `q/k/v/o_proj`.
@@ -76,10 +81,11 @@ Run commands through the local environment:
 UV_CACHE_DIR=.uv-cache uv run mlxrl --help
 ```
 
-Python 3.11 through 3.14 is supported. Runtime dependencies are intentionally
-small: `mlx` and `mlx-lm` on macOS arm64, matching the MLX wheel support window.
-Development dependencies include `pytest`, `ruff`, `pyright`, `mlx-tune`, and
-`mlx-lm-lora` for comparison benchmarks.
+Python 3.11+ is required. The MLX runtime path is supported on Apple Silicon
+with Python 3.11 through 3.14, matching the current MLX wheel support window.
+Runtime dependencies are intentionally small: `mlx` and `mlx-lm` on macOS
+arm64. Development dependencies include `pytest`, `ruff`, `pyright`,
+`mlx-tune`, and `mlx-lm-lora` for comparison benchmarks.
 The PyPI distribution name is `inductive-mlxrl`; the import package and console
 script remain `mlxrl`.
 
@@ -133,6 +139,14 @@ UV_CACHE_DIR=.uv-cache uv run mlxrl train \
   --available-memory-gb 48
 ```
 
+Reference multi-turn GiGPO run against the built-in recurring-state env:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl train \
+  --config examples/qwen3_0_6b_gigpo_recurring.toml \
+  --available-memory-gb 48
+```
+
 The config schema validates model id, quant bits, group size, completion/prompt
 lengths, checkpointing granularity, `iogpu.wired_limit_mb`, optimizer settings,
 algorithm hyperparameters, KL beta, and seed before a model is loaded. CLI
@@ -217,11 +231,71 @@ metrics = optimizer_step(
 )
 ```
 
+Multi-turn GiGPO uses the additive trajectory path:
+
+```python
+from mlxrl.algo import GiGPOAlgorithm
+from mlxrl.env import RecurringStateTextEnv
+from mlxrl.rollout import generate_agentic_trajectories
+from mlxrl.train import batch_from_trajectories, optimizer_step_trajectory
+
+algorithm = GiGPOAlgorithm(omega=1.0, gamma=1.0)
+trajectories = generate_agentic_trajectories(
+    model=model,
+    tokenizer=tokenizer,
+    env_factory=lambda task, seed, group: RecurringStateTextEnv(
+        task=str(task),
+        max_turns=3,
+    ),
+    tasks=["finish the task"],
+    group_size=4,
+    sampling=sampling,
+    rollout_mode="parallel_per_turn",
+)
+batch = batch_from_trajectories(
+    model=model,
+    trajectories=trajectories,
+    group_size=4,
+    pad_token_id=pad_token_id,
+    algorithm=algorithm,
+    compute_reference=beta != 0.0,
+)
+metrics = optimizer_step_trajectory(
+    model=model,
+    optimizer=optimizer,
+    batch=batch,
+    beta=beta,
+    pad_token_id=pad_token_id,
+    algorithm=algorithm,
+)
+```
+
 `micro_batch_size=0` keeps the original whole-batch path. Micro-batching is
 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.
@@ -252,6 +326,7 @@ import concrete algorithm implementations.
 | DAPO | asymmetric low/high clipping plus optional dynamic zero-advantage group filtering |
 | GSPO | sequence-level, length-normalized importance ratio and clipping |
 | RLOO | leave-one-out group baseline, no critic, no std-normalized advantage |
+| GiGPO | multi-turn episode advantage plus anchor-state step advantage |
 
 ## Memory Preflight
 
@@ -270,6 +345,9 @@ Qwen3.5-9B/G2/seq609/per-layer-checkpointed, `45.9 GB` for
 Qwen3.5-9B/G4/seq609/per-layer-checkpointed, and `36 GB` for
 Qwen3.5-9B/G2/seq128/no-checkpoint. For hybrid 9B no-checkpoint long-sequence
 configs, it reports an OOM-risk lower bound rather than a fake precise peak.
+Multi-turn configs include `max_turns` and `max_observation_len` in the
+estimated trajectory length and are labeled as estimated OOM-risk shapes until
+new measured anchors are added.
 For an obviously too-large Qwen3.5-9B/G8/prompt97/T512/no-checkpoint config on
 48 GB, it flags the run and suggests the measured-boundary fallback around
 G4/T512/checkpointed.
@@ -349,6 +427,7 @@ updates, or benchmark timing boundaries.
 ```text
 mlxrl/
   rollout/   # batched group generation
+  env/       # agentic environment protocols and reference envs
   policy/    # model loading, LoRA setup, logprob passes
   algo/      # GRPO-family advantages and losses
   train/     # value_and_grad and optimizer integration

inductive_mlxrl-0.2.2/examples/qwen35_9b_g2_checkpoint.toml

@@ -0,0 +1,29 @@
+model_id = "mlx-community/Qwen3.5-9B-MLX-4bit"
+quant_bits = 4
+group_size = 2
+max_prompt_len = 97
+max_completion_len = 512
+steps = 1
+seed = 7
+kl_beta = 0.0
+rank = 16
+scale = 2.0
+dropout = 0.0
+gradient_checkpointing = true
+checkpoint_granularity = "per-layer"
+iogpu_wired_limit_mb = 49152
+micro_batch_size = 0
+output = "reference_outputs/mlxrl_qwen35_9b_g2_checkpoint.npz"
+
+[optimizer]
+name = "adam"
+learning_rate = 0.00001
+
+[sampling]
+temperature = 0.7
+top_p = 0.95
+min_p = 0.0
+top_k = 0
+
+[algorithm]
+name = "grpo"

inductive_mlxrl-0.2.2/examples/qwen3_0_6b_gigpo_recurring.toml

@@ -0,0 +1,33 @@
+model_id = "mlx-community/Qwen3-0.6B-4bit"
+group_size = 4
+max_prompt_len = 512
+max_completion_len = 32
+max_turns = 3
+max_observation_len = 64
+rollout_mode = "parallel_per_turn"
+env_name = "recurring-text"
+steps = 5
+seed = 7
+kl_beta = 0.0
+rank = 8
+scale = 20.0
+dropout = 0.0
+gradient_checkpointing = false
+micro_batch_size = 0
+output = "reference_outputs/gigpo_recurring_reference.npz"
+
+[optimizer]
+name = "adam"
+learning_rate = 0.00001
+
+[sampling]
+temperature = 0.7
+top_p = 0.95
+min_p = 0.0
+top_k = 0
+
+[algorithm]
+name = "gigpo"
+gigpo_omega = 1.0
+gigpo_gamma = 1.0
+gigpo_normalization = "std"

inductive_mlxrl-0.2.2/examples/qwen3_0_6b_grpo.toml

@@ -0,0 +1,29 @@
+model_id = "mlx-community/Qwen3-0.6B-4bit"
+quant_bits = 4
+group_size = 2
+max_completion_len = 4
+max_prompt_len = 4096
+steps = 1
+seed = 7
+kl_beta = 0.0
+rank = 8
+scale = 20.0
+dropout = 0.0
+gradient_checkpointing = false
+checkpoint_granularity = "per-layer"
+iogpu_wired_limit_mb = 49152
+micro_batch_size = 0
+output = "reference_outputs/mlxrl_config_smoke.npz"
+
+[optimizer]
+name = "adam"
+learning_rate = 0.00001
+
+[sampling]
+temperature = 0.7
+top_p = 0.95
+min_p = 0.0
+top_k = 0
+
+[algorithm]
+name = "grpo"

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

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