inductive-mlxrl 0.1.2__tar.gz → 0.2.0__tar.gz

{inductive_mlxrl-0.1.2 → inductive_mlxrl-0.2.0}/CHANGELOG.md

@@ -5,6 +5,27 @@ 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.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

{inductive_mlxrl-0.1.2 → inductive_mlxrl-0.2.0}/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
@@ -105,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.2 → inductive_mlxrl-0.2.0}/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.2 → inductive_mlxrl-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inductive-mlxrl
-Version: 0.1.2
+Version: 0.2.0
 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
@@ -25,16 +25,18 @@ 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`.
@@ -159,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
@@ -243,6 +256,45 @@ 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
@@ -278,6 +330,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
 
@@ -296,6 +349,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.
@@ -375,6 +431,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.2 → inductive_mlxrl-0.2.0}/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`.
@@ -134,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
@@ -218,6 +231,45 @@ 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
@@ -253,6 +305,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
 
@@ -271,6 +324,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.
@@ -350,6 +406,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.0/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 = "/tmp/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.0/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.0/examples/qwen3_0_6b_grpo.toml

@@ -0,0 +1,30 @@
+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 = "/tmp/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.2 → inductive_mlxrl-0.2.0}/mlxrl/__init__.py

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

{inductive_mlxrl-0.1.2 → inductive_mlxrl-0.2.0}/mlxrl/algo/__init__.py

@@ -1,5 +1,12 @@
 """Algorithm-specific advantage and loss code."""
 
+from mlxrl.algo.gigpo import (
+    GiGPOAlgorithm,
+    anchor_state_groups,
+    anchor_state_step_advantages,
+    compute_gigpo_step_advantages,
+    flatten_step_records,
+)
 from mlxrl.algo.grpo import (
     AlgorithmLossMetrics,
     DAPOAlgorithm,
@@ -28,12 +35,17 @@ __all__ = [
     "GRPOAlgorithm",
     "GRPOLossMetrics",
     "GSPOAlgorithm",
+    "GiGPOAlgorithm",
     "PolicyAlgorithm",
     "RLOOAlgorithm",
     "algorithm_by_name",
+    "anchor_state_groups",
+    "anchor_state_step_advantages",
     "approximate_kl",
     "clip_ratio",
+    "compute_gigpo_step_advantages",
     "filter_zero_advantage_groups",
+    "flatten_step_records",
     "group_center_rewards",
     "group_normalize_rewards",
     "grpo_loss",