textpolicy 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

textpolicy/algorithms/grpo.py

@@ -26,7 +26,48 @@ if mx is None:
             return fn
 
     mx = _DummyMx()
-from typing import List, Union
+from typing import List, Union, Tuple, Dict, Any, Optional
+from dataclasses import dataclass
+from collections import defaultdict
+
+# Import length shaping utilities from dedicated module
+from .length_shaping import (
+    compute_length_penalty,
+    apply_length_shaping,
+    compute_length_shaping_stats,
+)
+
+
+# --- Clip Configuration Helper ---
+@dataclass
+class ClipConfig:
+    """Configuration for PPO/DAPO clipping bounds."""
+    low: float = 0.2
+    high: float = 0.28
+
+
+def resolve_clip_config(
+    clip_ratio: Optional[float],
+    clip_ratio_low: float = 0.2,
+    clip_ratio_high: float = 0.28,
+) -> ClipConfig:
+    """
+    Resolve clipping configuration with backward compatibility.
+
+    Centralizes the logic for handling symmetric vs asymmetric clipping bounds.
+
+    Args:
+        clip_ratio: Symmetric clipping ratio (backward compatibility).
+                   If provided, overrides clip_ratio_low and clip_ratio_high.
+        clip_ratio_low: Lower bound offset (default 0.2)
+        clip_ratio_high: Upper bound offset (default 0.28)
+
+    Returns:
+        ClipConfig with resolved low and high bounds
+    """
+    if clip_ratio is not None:
+        return ClipConfig(low=clip_ratio, high=clip_ratio)
+    return ClipConfig(low=clip_ratio_low, high=clip_ratio_high)
 
 
 def compute_advantages(rewards: Union[List[float], mx.array]) -> mx.array:
@@ -113,45 +154,88 @@ def compute_advantages_dr_grpo(rewards: Union[List[float], mx.array]) -> mx.arra
 
 def policy_loss(
     old_logprobs: mx.array,
-    new_logprobs: mx.array, 
+    new_logprobs: mx.array,
     advantages: mx.array,
-    clip_ratio: float = 0.2
+    clip_ratio: float = None,
+    clip_ratio_low: float = 0.2,
+    clip_ratio_high: float = 0.28,
+    normalize_constant: int = None
 ) -> mx.array:
     """
-    GRPO policy loss with PPO-style clipping.
-    
+    GRPO policy loss with PPO-style clipping, supporting DAPO asymmetric bounds.
+
     Uses clipped surrogate objective but with group-relative advantages
-    instead of GAE advantages.
-    
+    instead of GAE advantages. Supports asymmetric clipping bounds (DAPO-style)
+    to prevent entropy collapse while maintaining training stability.
+
+    DAPO insight: Asymmetric bounds allow the model to increase probabilities
+    of good actions more easily than decreasing probabilities of bad actions,
+    promoting diversity and preventing entropy collapse.
+
+    Dr. GRPO insight: Dividing by a fixed constant instead of token count
+    eliminates length bias that artificially inflates incorrect (longer) responses.
+
     Args:
         old_logprobs: Log probabilities from rollout collection
         new_logprobs: Log probabilities from current policy evaluation
         advantages: Group-relative advantages from compute_advantages()
-        clip_ratio: Clipping ratio for surrogate objective
-        
+        clip_ratio: Symmetric clipping ratio (for backward compatibility).
+                   If provided, overrides clip_ratio_low and clip_ratio_high.
+        clip_ratio_low: Lower bound offset (default 0.2, gives lower bound of 0.8)
+        clip_ratio_high: Upper bound offset (default 0.28, gives upper bound of 1.28)
+        normalize_constant: Fixed constant divisor for loss normalization.
+                           If None (default), uses mean (original behavior).
+                           If provided, uses sum/constant to eliminate length bias.
+                           Typical values: 1024, or batch_size.
+
     Returns:
         Policy loss scalar (to be minimized)
-        
+
     Notes:
     - Fully vectorized (no Python loops over batch)
     - Uses in-place operations where possible
     - Suitable for MLX graph optimization
     - Single forward pass through computation
+    - DAPO defaults: clip_ratio_low=0.2, clip_ratio_high=0.28
+    - Length bias: When using mean, longer sequences have lower per-token
+      contribution, creating implicit bias toward short responses.
+
+    References:
+        DAPO: An Open-Source LLM Reinforcement Learning System at Scale
+        https://arxiv.org/abs/2503.14476
+
+        Dr. GRPO: Understanding R1-Zero-Like Training
+        https://arxiv.org/abs/2503.20783
     """
-    # Importance ratio: π_new / π_old  
+    # Resolve clipping configuration (handles backward compatibility)
+    clip_cfg = resolve_clip_config(clip_ratio, clip_ratio_low, clip_ratio_high)
+
+    # Importance ratio: π_new / π_old
     # MLX optimizes exp() for Apple Silicon
     ratio = mx.exp(new_logprobs - old_logprobs)
-    
-    # PPO clipped surrogate objective
-    # L = min(ratio * A, clip(ratio, 1-ε, 1+ε) * A)
-    clipped_ratio = mx.clip(ratio, 1 - clip_ratio, 1 + clip_ratio)
-    
-    # Element-wise minimum and mean reduction
-    # Negative because we minimize (original maximizes)
+
+    # PPO clipped surrogate objective with asymmetric bounds (DAPO-style)
+    # L = min(ratio * A, clip(ratio, 1-ε_low, 1+ε_high) * A)
+    clipped_ratio = mx.clip(ratio, 1 - clip_cfg.low, 1 + clip_cfg.high)
+
+    # Element-wise minimum
     surr1 = ratio * advantages
     surr2 = clipped_ratio * advantages
-    loss = -mx.mean(mx.minimum(surr1, surr2))
-    
+    min_surr = mx.minimum(surr1, surr2)
+
+    # Normalization: either mean (original) or sum/constant (Dr. GRPO)
+    if normalize_constant is not None:
+        if normalize_constant <= 0:
+            raise ValueError(
+                f"normalize_constant must be positive, got {normalize_constant}"
+            )
+        # Fixed constant normalization eliminates length bias
+        # All sequences contribute equally regardless of length
+        loss = -mx.sum(min_surr) / normalize_constant
+    else:
+        # Original mean behavior (for backward compatibility)
+        loss = -mx.mean(min_surr)
+
     return loss
 
 
@@ -163,86 +247,491 @@ def compute_advantages_compiled(rewards: mx.array) -> mx.array:
     return rewards - group_mean
 
 
-@mx.compile  
-def policy_loss_compiled(
+# --- Compiled Policy Loss Variants ---
+# Two internal compiled functions for different normalization strategies.
+# Compiled functions require static control flow, so we keep them separate.
+
+@mx.compile
+def _policy_loss_compiled_mean(
     old_logprobs: mx.array,
     new_logprobs: mx.array,
     advantages: mx.array,
-    clip_ratio: float = 0.2
+    clip_ratio_low: float,
+    clip_ratio_high: float
 ) -> mx.array:
-    """Compiled version of policy_loss for maximum performance."""
+    """Internal compiled function: mean normalization."""
     ratio = mx.exp(new_logprobs - old_logprobs)
-    clipped_ratio = mx.clip(ratio, 1 - clip_ratio, 1 + clip_ratio)
+    clipped_ratio = mx.clip(ratio, 1 - clip_ratio_low, 1 + clip_ratio_high)
     surr1 = ratio * advantages
     surr2 = clipped_ratio * advantages
     return -mx.mean(mx.minimum(surr1, surr2))
 
 
+@mx.compile
+def _policy_loss_compiled_constant(
+    old_logprobs: mx.array,
+    new_logprobs: mx.array,
+    advantages: mx.array,
+    clip_ratio_low: float,
+    clip_ratio_high: float,
+    normalize_constant: float
+) -> mx.array:
+    """Internal compiled function: constant normalization."""
+    ratio = mx.exp(new_logprobs - old_logprobs)
+    clipped_ratio = mx.clip(ratio, 1 - clip_ratio_low, 1 + clip_ratio_high)
+    surr1 = ratio * advantages
+    surr2 = clipped_ratio * advantages
+    return -mx.sum(mx.minimum(surr1, surr2)) / normalize_constant
+
+
+def policy_loss_compiled(
+    old_logprobs: mx.array,
+    new_logprobs: mx.array,
+    advantages: mx.array,
+    clip_ratio: float = None,
+    clip_ratio_low: float = 0.2,
+    clip_ratio_high: float = 0.28
+) -> mx.array:
+    """
+    Compiled version of policy_loss for maximum performance (mean normalization).
+
+    Supports DAPO-style asymmetric clipping bounds with backward compatibility.
+    Uses mean normalization (original behavior).
+
+    Args:
+        old_logprobs: Log probabilities from rollout collection
+        new_logprobs: Log probabilities from current policy evaluation
+        advantages: Group-relative advantages
+        clip_ratio: Symmetric clipping ratio (for backward compatibility).
+                   If provided, overrides clip_ratio_low and clip_ratio_high.
+        clip_ratio_low: Lower bound offset (default 0.2)
+        clip_ratio_high: Upper bound offset (default 0.28)
+    """
+    clip_cfg = resolve_clip_config(clip_ratio, clip_ratio_low, clip_ratio_high)
+    return _policy_loss_compiled_mean(
+        old_logprobs, new_logprobs, advantages,
+        clip_cfg.low, clip_cfg.high
+    )
+
+
+def policy_loss_compiled_constant_norm(
+    old_logprobs: mx.array,
+    new_logprobs: mx.array,
+    advantages: mx.array,
+    clip_ratio: float = None,
+    clip_ratio_low: float = 0.2,
+    clip_ratio_high: float = 0.28,
+    normalize_constant: float = 1024.0
+) -> mx.array:
+    """
+    Compiled version of policy_loss with fixed constant normalization (Dr. GRPO).
+
+    Uses sum/constant instead of mean to eliminate length bias.
+
+    Args:
+        old_logprobs: Log probabilities from rollout collection
+        new_logprobs: Log probabilities from current policy evaluation
+        advantages: Group-relative advantages
+        clip_ratio: Symmetric clipping ratio (for backward compatibility).
+                   If provided, overrides clip_ratio_low and clip_ratio_high.
+        clip_ratio_low: Lower bound offset (default 0.2)
+        clip_ratio_high: Upper bound offset (default 0.28)
+        normalize_constant: Fixed constant divisor (default 1024)
+
+    References:
+        Dr. GRPO: Understanding R1-Zero-Like Training
+        https://arxiv.org/abs/2503.20783
+    """
+    if normalize_constant <= 0:
+        raise ValueError(
+            f"normalize_constant must be positive, got {normalize_constant}"
+        )
+    clip_cfg = resolve_clip_config(clip_ratio, clip_ratio_low, clip_ratio_high)
+    return _policy_loss_compiled_constant(
+        old_logprobs, new_logprobs, advantages,
+        clip_cfg.low, clip_cfg.high, normalize_constant
+    )
+
+
 def entropy_bonus(logprobs: mx.array, coefficient: float = 0.01) -> mx.array:
     """
     Entropy bonus for exploration (optional GRPO component).
-    
+
     Args:
         logprobs: Log probabilities from policy
         coefficient: Entropy coefficient (typically small, like 0.01)
-        
+
     Returns:
         Entropy bonus (added to loss for exploration)
     """
     if coefficient <= 0:
         return mx.array(0.0)
-    
+
     # Entropy = -sum(p * log(p))
     # For log probabilities: entropy = -sum(exp(logp) * logp)
     probs = mx.exp(logprobs)
     entropy = -mx.sum(probs * logprobs, axis=-1)
-    
+
     # Return negative entropy (since we add to loss but want to maximize entropy)
     return -coefficient * mx.mean(entropy)
 
 
+# Note: compute_length_penalty, apply_length_shaping, and compute_length_shaping_stats
+# are imported from .length_shaping module (see imports at top of file)
+
+
+# DAPO-style dynamic batch filtering (Issue #9)
+def _get_episode_reward(episode) -> float:
+    """Extract total reward from episode (handles both Episode objects and dicts)."""
+    if hasattr(episode, 'rew'):
+        # Episode object
+        return float(mx.sum(mx.array(episode.rew)).item())
+    else:
+        # Serialized dictionary
+        rew = episode.get('rew', episode.get('reward', [0.0]))
+        if isinstance(rew, (int, float)):
+            return float(rew)
+        return float(mx.sum(mx.array(rew)).item())
+
+
+def _get_prompt_key(episode) -> tuple:
+    """
+    Generate a hashable key for an episode's prompt.
+
+    Handles both Episode objects and serialized dictionaries.
+    Uses the observation (prompt) tokens to identify the prompt.
+    """
+    if hasattr(episode, 'obs'):
+        obs = episode.obs
+    else:
+        obs = episode.get('obs', [])
+
+    # Flatten nested structures to create consistent key
+    flattened = []
+    for item in obs:
+        if hasattr(item, 'tolist'):  # MLX array
+            flattened.extend(item.tolist())
+        elif isinstance(item, list):
+            flattened.extend(item)
+        else:
+            flattened.append(item)
+
+    return tuple(flattened)
+
+
+def _precompute_episode_rewards(episodes: List[Any]) -> List[float]:
+    """
+    Pre-compute rewards for all episodes in a single pass.
+
+    Uses batched MLX evaluation to avoid per-episode .item() sync barriers.
+    All mx.sum() calls are built lazily, then evaluated in one mx.eval() call.
+
+    Args:
+        episodes: List of episodes
+
+    Returns:
+        List of rewards in the same order as episodes
+    """
+    if not episodes:
+        return []
+
+    rewards: List[Optional[float]] = [None] * len(episodes)
+    pending: List[Tuple[int, mx.array]] = []  # (index, lazy_sum) pairs
+
+    for i, ep in enumerate(episodes):
+        if hasattr(ep, 'rew'):
+            rew = ep.rew
+        else:
+            rew = ep.get('rew', ep.get('reward', [0.0]))
+
+        if isinstance(rew, (int, float)):
+            rewards[i] = float(rew)
+        else:
+            pending.append((i, mx.sum(mx.array(rew))))
+
+    # Single sync barrier for all array rewards
+    if pending:
+        indices, lazy_sums = zip(*pending)
+        stacked = mx.stack(list(lazy_sums))
+        mx.eval(stacked)
+        values = stacked.tolist()
+        for idx, val in zip(indices, values):
+            rewards[idx] = float(val)
+
+    return rewards  # type: ignore[return-value]
+
+
+def _compute_group_variance_and_mean(
+    group_indices: List[int],
+    all_rewards: List[float]
+) -> Tuple[float, float]:
+    """
+    Compute variance and mean for a group of episodes using pre-computed rewards.
+
+    Args:
+        group_indices: Indices into all_rewards for this group
+        all_rewards: Pre-computed rewards for all episodes
+
+    Returns:
+        Tuple of (variance, mean)
+    """
+    group_rewards = mx.array([all_rewards[i] for i in group_indices])
+    return mx.var(group_rewards).item(), mx.mean(group_rewards).item()
+
+
+def filter_informative_prompts(
+    episodes: List[Any],
+    min_variance: float = 0.01,
+    keep_single_completion: bool = True
+) -> Tuple[List[Any], Dict[str, Union[int, float]]]:
+    """
+    Filter episodes to keep only informative prompts (DAPO dynamic sampling).
+
+    Removes prompts where all completions have same outcome:
+    - All correct (reward ~1.0): no learning signal (nothing to improve)
+    - All wrong (reward ~0.0): no positive signal (can't learn what works)
+
+    GRPO uses group-relative advantages. If all completions have the same
+    outcome, advantages are zero, producing no gradient and wasting compute.
+
+    Note on single-completion prompts:
+        The DAPO paper (Equation 11) defines informative prompts as having
+        mixed outcomes: `0 < |correct| < G`. This assumes G > 1 completions
+        per prompt. For single-completion prompts (G=1), variance is always 0
+        by definition, but this doesn't mean "all outcomes are the same" -
+        it means we have insufficient data to determine variance.
+
+        By default (keep_single_completion=True), single-completion prompts
+        are kept since they still provide valid gradient signal. Set to False
+        to filter them out (stricter DAPO interpretation).
+
+    Args:
+        episodes: List of episodes (Episode objects or serialized dicts)
+        min_variance: Minimum reward variance to keep a prompt group.
+                     Groups with variance below this threshold are filtered out.
+                     Default 0.01 filters prompts with essentially identical rewards.
+                     Only applied to groups with 2+ completions.
+        keep_single_completion: Whether to keep prompts with only one completion.
+                               Default True (keep them). Set False to require
+                               multiple completions for variance calculation.
+
+    Returns:
+        Tuple of:
+        - filtered: List of episodes from informative prompts
+        - stats: Dictionary with filtering statistics:
+            - 'prompts_kept': Number of prompt groups kept
+            - 'prompts_dropped_all_correct': Prompts where all completions succeeded
+            - 'prompts_dropped_all_wrong': Prompts where all completions failed
+            - 'prompts_dropped_single': Prompts dropped due to single completion
+            - 'prompts_kept_single': Single-completion prompts that were kept
+            - 'episodes_kept': Total episodes kept
+            - 'episodes_dropped': Total episodes filtered out
+            - 'filter_rate': Fraction of prompts filtered
+
+    Example:
+        >>> filtered, stats = filter_informative_prompts(episodes, min_variance=0.01)
+        >>> print(f"Kept {stats['prompts_kept']} prompts, "
+        ...       f"dropped {stats['prompts_dropped_all_correct']} all-correct, "
+        ...       f"{stats['prompts_dropped_all_wrong']} all-wrong")
+
+    References:
+        DAPO: An Open-Source LLM Reinforcement Learning System at Scale
+        https://arxiv.org/abs/2503.14476 (Equation 11: 0 < |correct| < G)
+
+        GRPO++ Tricks
+        https://cameronrwolfe.substack.com/p/grpo-tricks
+    """
+    if not episodes:
+        return [], {
+            'prompts_kept': 0,
+            'prompts_dropped_all_correct': 0,
+            'prompts_dropped_all_wrong': 0,
+            'prompts_dropped_single': 0,
+            'prompts_kept_single': 0,
+            'episodes_kept': 0,
+            'episodes_dropped': 0,
+            'filter_rate': 0.0,
+        }
+
+    # Pre-compute all rewards once (avoids repeated _get_episode_reward calls)
+    all_rewards = _precompute_episode_rewards(episodes)
+
+    # Group episodes by prompt, storing indices instead of episodes
+    prompt_groups: Dict[tuple, List[int]] = defaultdict(list)
+    for idx, ep in enumerate(episodes):
+        prompt_key = _get_prompt_key(ep)
+        prompt_groups[prompt_key].append(idx)
+
+    filtered = []
+    stats = {
+        'prompts_kept': 0,
+        'prompts_dropped_all_correct': 0,
+        'prompts_dropped_all_wrong': 0,
+        'prompts_dropped_single': 0,
+        'prompts_kept_single': 0,
+        'episodes_kept': 0,
+        'episodes_dropped': 0,
+    }
+
+    for prompt_key, group_indices in prompt_groups.items():
+        group_size = len(group_indices)
+
+        # Handle single-completion prompts separately
+        if group_size == 1:
+            if keep_single_completion:
+                # Keep single-completion prompts (variance undefined, not "zero")
+                filtered.append(episodes[group_indices[0]])
+                stats['prompts_kept'] += 1
+                stats['prompts_kept_single'] += 1
+                stats['episodes_kept'] += 1
+            else:
+                # Filter out single-completion prompts (strict DAPO interpretation)
+                stats['prompts_dropped_single'] += 1
+                stats['episodes_dropped'] += 1
+            continue
+
+        # For groups with 2+ completions, use variance criterion
+        variance, mean_reward = _compute_group_variance_and_mean(group_indices, all_rewards)
+
+        if variance > min_variance:
+            # Informative: mixed outcomes, keep all episodes from this prompt
+            for idx in group_indices:
+                filtered.append(episodes[idx])
+            stats['prompts_kept'] += 1
+            stats['episodes_kept'] += group_size
+        else:
+            # Uninformative: all completions have same outcome
+            stats['episodes_dropped'] += group_size
+            if mean_reward > 0.5:
+                stats['prompts_dropped_all_correct'] += 1
+            else:
+                stats['prompts_dropped_all_wrong'] += 1
+
+    # Compute filter rate
+    total_prompts = len(prompt_groups)
+    stats['filter_rate'] = 1.0 - (stats['prompts_kept'] / total_prompts) if total_prompts > 0 else 0.0
+
+    return filtered, stats
+
+
+def compute_prompt_group_stats(episodes: List[Any]) -> Dict[str, Any]:
+    """
+    Compute statistics about prompt groups for monitoring.
+
+    Useful for understanding the distribution of prompts and completions
+    before and after filtering.
+
+    Args:
+        episodes: List of episodes
+
+    Returns:
+        Dictionary with:
+        - 'num_prompts': Total unique prompts
+        - 'num_episodes': Total episodes
+        - 'completions_per_prompt': Average completions per prompt
+        - 'reward_variance_mean': Mean variance across prompt groups
+        - 'reward_variance_std': Std of variance across prompt groups
+    """
+    if not episodes:
+        return {
+            'num_prompts': 0,
+            'num_episodes': 0,
+            'completions_per_prompt': 0.0,
+            'reward_variance_mean': 0.0,
+            'reward_variance_std': 0.0,
+        }
+
+    # Pre-compute all rewards once
+    all_rewards = _precompute_episode_rewards(episodes)
+
+    # Group by prompt, storing indices
+    prompt_groups: Dict[tuple, List[int]] = defaultdict(list)
+    for idx, ep in enumerate(episodes):
+        prompt_key = _get_prompt_key(ep)
+        prompt_groups[prompt_key].append(idx)
+
+    # Compute variance for each group using pre-computed rewards
+    variances = []
+    for group_indices in prompt_groups.values():
+        group_rewards = mx.array([all_rewards[i] for i in group_indices])
+        variances.append(mx.var(group_rewards).item())
+
+    variances_arr = mx.array(variances) if variances else mx.array([0.0])
+
+    return {
+        'num_prompts': len(prompt_groups),
+        'num_episodes': len(episodes),
+        'completions_per_prompt': len(episodes) / len(prompt_groups) if prompt_groups else 0.0,
+        'reward_variance_mean': float(mx.mean(variances_arr).item()),
+        'reward_variance_std': float(mx.std(variances_arr).item()),
+    }
+
+
 # Convenience function for complete GRPO computation
 def grpo_loss(
     old_logprobs: mx.array,
     new_logprobs: mx.array,
     rewards: Union[List[float], mx.array],
-    clip_ratio: float = 0.2,
-    entropy_coeff: float = 0.0
+    clip_ratio: float = None,
+    clip_ratio_low: float = 0.2,
+    clip_ratio_high: float = 0.28,
+    entropy_coeff: float = 0.0,
+    normalize_constant: int = None
 ) -> mx.array:
     """
     Complete GRPO loss computation in one function.
-    
+
     Combines advantage calculation and policy loss for convenience.
     Can be compiled as a single unit for maximum efficiency.
-    
+    Supports DAPO-style asymmetric clipping bounds and Dr. GRPO length-bias fix.
+
     Args:
         old_logprobs: Log probabilities from rollout
         new_logprobs: Log probabilities from current policy
         rewards: Episode rewards for group-relative advantages
-        clip_ratio: PPO clipping ratio
+        clip_ratio: Symmetric clipping ratio (for backward compatibility).
+                   If provided, overrides clip_ratio_low and clip_ratio_high.
+        clip_ratio_low: Lower bound offset (default 0.2)
+        clip_ratio_high: Upper bound offset (default 0.28)
         entropy_coeff: Entropy bonus coefficient (0 disables)
-        
+        normalize_constant: Fixed constant divisor for loss normalization.
+                           If None (default), uses mean. If provided, uses
+                           sum/constant to eliminate length bias.
+
     Returns:
         Total GRPO loss (policy + optional entropy)
+
+    References:
+        DAPO: An Open-Source LLM Reinforcement Learning System at Scale
+        https://arxiv.org/abs/2503.14476
+
+        Dr. GRPO: Understanding R1-Zero-Like Training
+        https://arxiv.org/abs/2503.20783
     """
     # Compute group-relative advantages
     advantages = compute_advantages(rewards)
-    
+
     # Expand advantages to match logprob sequence length if needed
     if advantages.ndim == 1 and old_logprobs.ndim > 1:
         # Each episode contributes its advantage to all tokens in that episode
         # This requires knowing episode boundaries - simplified version assumes
         # advantages and logprobs are already aligned
         pass
-    
-    # Compute policy loss
-    policy_loss_val = policy_loss(old_logprobs, new_logprobs, advantages, clip_ratio)
-    
+
+    # Compute policy loss with asymmetric clipping and optional length-bias fix
+    policy_loss_val = policy_loss(
+        old_logprobs, new_logprobs, advantages,
+        clip_ratio=clip_ratio,
+        clip_ratio_low=clip_ratio_low,
+        clip_ratio_high=clip_ratio_high,
+        normalize_constant=normalize_constant
+    )
+
     # Add entropy bonus if specified
     if entropy_coeff > 0:
         entropy_bonus_val = entropy_bonus(new_logprobs, entropy_coeff)
         return policy_loss_val + entropy_bonus_val
-    
+
     return policy_loss_val
 
 
@@ -251,392 +740,313 @@ def compute_metrics(
     old_logprobs: mx.array,
     new_logprobs: mx.array,
     advantages: mx.array,
-    clip_ratio: float = 0.2
+    clip_ratio: float = None,
+    clip_ratio_low: float = 0.2,
+    clip_ratio_high: float = 0.28
 ) -> dict:
     """
     Compute GRPO training metrics for monitoring.
-    
+
+    Supports DAPO-style asymmetric clipping bounds and tracks clip fractions
+    for upper vs lower bounds separately.
+
     Args:
         old_logprobs: Log probabilities from rollout
-        new_logprobs: Log probabilities from current policy  
+        new_logprobs: Log probabilities from current policy
         advantages: Group-relative advantages
-        clip_ratio: Clipping ratio used in loss
-        
+        clip_ratio: Symmetric clipping ratio (for backward compatibility).
+                   If provided, overrides clip_ratio_low and clip_ratio_high.
+        clip_ratio_low: Lower bound offset (default 0.2)
+        clip_ratio_high: Upper bound offset (default 0.28)
+
     Returns:
-        Dictionary of metrics for logging/monitoring
+        Dictionary of metrics for logging/monitoring, including:
+        - clip_fraction_lower: Fraction of ratios clipped at lower bound
+        - clip_fraction_upper: Fraction of ratios clipped at upper bound
+        - clip_fraction: Total fraction of ratios clipped (either bound)
     """
+    # Resolve clipping configuration (handles backward compatibility)
+    clip_cfg = resolve_clip_config(clip_ratio, clip_ratio_low, clip_ratio_high)
+
     # Importance ratio statistics
     ratio = mx.exp(new_logprobs - old_logprobs)
-    
-    # Clipping statistics
-    clip_lower = 1 - clip_ratio
-    clip_upper = 1 + clip_ratio
-    clipped = (ratio < clip_lower) | (ratio > clip_upper)
+
+    # Asymmetric clipping bounds
+    clip_lower = 1 - clip_cfg.low
+    clip_upper = 1 + clip_cfg.high
+
+    # Track clip fractions separately for upper and lower bounds
+    clipped_lower = ratio < clip_lower
+    clipped_upper = ratio > clip_upper
+    clipped = clipped_lower | clipped_upper
+
+    clip_fraction_lower = mx.mean(clipped_lower.astype(mx.float32))
+    clip_fraction_upper = mx.mean(clipped_upper.astype(mx.float32))
     clip_fraction = mx.mean(clipped.astype(mx.float32))
-    
+
     # KL divergence approximation
     kl_div = mx.mean(old_logprobs - new_logprobs)
-    
+
     return {
         'mean_advantage': mx.mean(advantages).item(),
         'std_advantage': mx.std(advantages).item(),
         'mean_ratio': mx.mean(ratio).item(),
         'clip_fraction': clip_fraction.item(),
+        'clip_fraction_lower': clip_fraction_lower.item(),
+        'clip_fraction_upper': clip_fraction_upper.item(),
         'kl_divergence': kl_div.item(),
         'min_advantage': mx.min(advantages).item(),
-        'max_advantage': mx.max(advantages).item()
+        'max_advantage': mx.max(advantages).item(),
+        'clip_ratio_low': clip_cfg.low,
+        'clip_ratio_high': clip_cfg.high
     }
 
 
-# Algorithm-specific data selection strategies
-def select_all_data(buffer):
+# --- Episode Packing Helper ---
+def _flatten_tokens(items: List[Any]) -> List:
+    """Flatten nested token sequences into a flat list."""
+    flattened = []
+    for item in items:
+        if hasattr(item, 'tolist'):  # MLX array
+            flattened.extend(item.tolist())
+        elif isinstance(item, list):  # Python list
+            flattened.extend(item)
+        else:  # Single token
+            flattened.append(item)
+    return flattened
+
+
+def _pack_episodes(episodes: List[Any]) -> Dict[str, Any]:
     """
-    GRPO data selector: Use all available data.
-    
-    GRPO is on-policy but can benefit from using all collected episodes
-    since group-relative advantages normalize across the entire group.
-    
+    Pack episodes into batch data for GRPO training.
+
+    This is the shared helper for episode-to-batch conversion, used by all
+    data selectors (select_all_data, select_informative_data, select_recent_data).
+
     Args:
-        buffer: Buffer containing episodes
-        
+        episodes: List of episodes (Episode objects or serialized dicts)
+
     Returns:
-        All episode data prepared for training
+        Dictionary with:
+        - 'obs': Flat concatenated full sequences (prompt + response)
+        - 'act': Flat concatenated response tokens
+        - 'logprob': Flat concatenated log probabilities
+        - 'rewards': Episode rewards as MLX array
+        - 'episode_lengths': List of episode lengths
     """
-    from textpolicy.buffer import Buffer
-    if not isinstance(buffer, Buffer):
-        raise TypeError(f"Expected Buffer, got {type(buffer)}")
-    
-    # Use all available data - GRPO benefits from larger groups
-    episodes_data = buffer.sample()  # This returns concatenated transitions
-    
-    # We need to convert this back to episode structure for reward extraction
-    episodes = buffer.episodes  # Access episodes directly from storage
-    
     if not episodes:
-        raise ValueError("Buffer is empty - no episodes to train on")
-    
-    # Extract episode rewards for advantage computation
-    episode_rewards = []
+        return {
+            'obs': mx.array([], dtype=mx.int64),
+            'act': mx.array([], dtype=mx.int64),
+            'logprob': mx.array([], dtype=mx.float32),
+            'rewards': mx.array([]),
+            'episode_lengths': [],
+        }
+
     episode_lengths = []
-    
-    # Collect all transitions
     all_obs = []
     all_acts = []
     all_logprobs = []
-    
-    for episode in episodes:
-        # Episode reward (sum of all rewards in episode)
-        # Handle both Episode objects and serialized dictionaries
+    pending_reward_sums: List[Tuple[int, mx.array]] = []
+    scalar_rewards: Dict[int, float] = {}
+
+    for i, episode in enumerate(episodes):
         if hasattr(episode, 'rew'):
             # Episode object with attributes
-            episode_reward = mx.sum(mx.array(episode.rew)).item()
-            episode_rewards.append(episode_reward)
-            episode_lengths.append(len(episode.obs))
-            
-            # Collect transitions
-            # For proper logprob extraction during training, we need the full context (prompt + response)
-            # This matches how the model was called during rollout generation
-            # Flatten nested token sequences to create uniform token arrays
-            
-            # Extract and flatten observation tokens (prompt)
-            flattened_obs = []
-            for obs in episode.obs:
-                if hasattr(obs, 'tolist'):  # MLX array
-                    flattened_obs.extend(obs.tolist())
-                elif isinstance(obs, list):  # Python list
-                    flattened_obs.extend(obs)
-                else:  # Single token
-                    flattened_obs.append(obs)
-            
-            # Extract and flatten action tokens (response)
-            flattened_acts = []
-            for act in episode.act:
-                if hasattr(act, 'tolist'):  # MLX array
-                    flattened_acts.extend(act.tolist())
-                elif isinstance(act, list):  # Python list
-                    flattened_acts.extend(act)
-                else:  # Single token
-                    flattened_acts.append(act)
-            
+            pending_reward_sums.append((i, mx.sum(mx.array(episode.rew))))
+
+            # Flatten observation and action tokens
+            flattened_obs = _flatten_tokens(episode.obs)
+            flattened_acts = _flatten_tokens(episode.act)
+
+            # Use flattened token count for episode_lengths (used by _expand_advantages)
+            # This ensures alignment between expanded advantages and actual token sequences
+            episode_lengths.append(len(flattened_acts))
+
             # Create full sequence: [prompt_tokens..., response_tokens...]
             full_sequence = flattened_obs + flattened_acts
             all_obs.append(full_sequence)
             all_acts.append(flattened_acts)
-            all_logprobs.append(episode.logprob if episode.logprob else [])
+            all_logprobs.append(episode.logprob if episode.logprob is not None else [])
         else:
             # Serialized dictionary from multiprocessing
-            episode_reward = mx.sum(episode['rew']).item()
-            episode_rewards.append(episode_reward)
-            episode_lengths.append(len(episode['obs']))
-            
-            # Collect transitions
-            # For proper logprob extraction during training, we need the full context (prompt + response)
-            # This matches how the model was called during rollout generation
-            full_sequence = episode['obs'] + episode['act']  # Concatenate prompt + response
+            rew = episode['rew']
+            if isinstance(rew, (int, float)):
+                scalar_rewards[i] = float(rew)
+            else:
+                pending_reward_sums.append((i, mx.sum(mx.array(rew))))
+
+            # Flatten observation and action tokens
+            flattened_obs = _flatten_tokens(episode['obs'])
+            flattened_acts = _flatten_tokens(episode['act'])
+
+            # Use flattened token count for episode_lengths
+            episode_lengths.append(len(flattened_acts))
+
+            full_sequence = flattened_obs + flattened_acts
             all_obs.append(full_sequence)
-            all_acts.append(episode['act'])
+            all_acts.append(flattened_acts)
             all_logprobs.append(episode.get('logprob', []))
-    
-    # Convert Python lists to MLX arrays before concatenation
-    # This is required because Episode objects store data as Python lists for memory efficiency
-    # For proper logprob extraction, we need uniform-length sequences, so we pad to the maximum length
-    
-    # Find maximum sequence length for padding
+
+    # Batch evaluate all pending reward sums (single sync barrier instead of N)
+    episode_rewards = [0.0] * len(episodes)
+    for idx, val in scalar_rewards.items():
+        episode_rewards[idx] = val
+    if pending_reward_sums:
+        indices, lazy_sums = zip(*pending_reward_sums)
+        stacked = mx.stack(list(lazy_sums))
+        mx.eval(stacked)
+        values = stacked.tolist()
+        for idx, val in zip(indices, values):
+            episode_rewards[idx] = float(val)
+
+    # Find maximum sequence lengths for padding
     max_obs_len = max(len(obs) for obs in all_obs) if all_obs else 0
     max_act_len = max(len(act) for act in all_acts) if all_acts else 0
     max_logprob_len = max(len(logprob) for logprob in all_logprobs) if all_logprobs else 0
-    
-    # MLX-native padding and array operations for optimal Apple Silicon performance
-    # Convert all sequences to MLX arrays and pad directly in MLX space
-    try:
-        # Convert all sequences to MLX arrays first (staying in unified memory)
-        all_obs_mx = [mx.array(obs, dtype=mx.int64) for obs in all_obs if obs]
-        all_acts_mx = [mx.array(act, dtype=mx.int64) for act in all_acts if act]
-        all_logprobs_mx = [mx.array(logprob, dtype=mx.float32) for logprob in all_logprobs if logprob]
-        
-        # Pad using native MLX operations (more efficient for Apple Silicon)
-        if all_obs_mx:
-            padded_obs_mx = [mx.pad(obs, (0, max_obs_len - obs.shape[0]), constant_values=0) 
-                           if obs.shape[0] < max_obs_len else obs[:max_obs_len] 
-                           for obs in all_obs_mx]
-        else:
-            padded_obs_mx = []
-            
-        if all_acts_mx:
-            padded_acts_mx = [mx.pad(act, (0, max_act_len - act.shape[0]), constant_values=0) 
-                            if act.shape[0] < max_act_len else act[:max_act_len] 
-                            for act in all_acts_mx]
-        else:
-            padded_acts_mx = []
-            
-        if all_logprobs_mx:
-            padded_logprobs_mx = [mx.pad(logprob, (0, max_logprob_len - logprob.shape[0]), constant_values=0.0) 
-                                if logprob.shape[0] < max_logprob_len else logprob[:max_logprob_len] 
-                                for logprob in all_logprobs_mx]
-        else:
-            padded_logprobs_mx = []
-        
-        # Use padded MLX arrays directly (no intermediate conversion needed)
-        all_obs_mx = padded_obs_mx
-        all_acts_mx = padded_acts_mx  
-        all_logprobs_mx = padded_logprobs_mx
-        
-    except Exception as e:
-        print(f"ERROR in MLX array conversion: {e}")
-        print(f"DEBUG: all_obs types: {[type(obs) for obs in all_obs[:3]]}")  # Show first 3 for brevity
-        print(f"DEBUG: all_logprobs types: {[type(logprob) for logprob in all_logprobs[:3]]}")
-        raise
-    
-    # GRPO data structure: both observations and actions as flat concatenated sequences
-    # This matches the expected format for GRPO logprob extraction function
-    batch_data = {
-        'obs': mx.concatenate(all_obs_mx) if all_obs_mx else mx.array([]),  # Flat concatenated full sequences
-        'act': mx.concatenate(all_acts_mx) if all_acts_mx else mx.array([]),  # Flat concatenated response tokens
-        'logprob': mx.concatenate([logprob.flatten() for logprob in all_logprobs_mx]) if all_logprobs_mx else mx.array([]),  # Flat sequence for training
+
+    # Convert to MLX arrays with padding
+    # Always create an array for each episode to maintain alignment
+    all_obs_mx = [mx.array(obs, dtype=mx.int64) if obs else mx.array([], dtype=mx.int64) for obs in all_obs]
+    all_acts_mx = [mx.array(act, dtype=mx.int64) if act else mx.array([], dtype=mx.int64) for act in all_acts]
+    all_logprobs_mx = [mx.array(logprob, dtype=mx.float32) if logprob else mx.array([], dtype=mx.float32) for logprob in all_logprobs]
+
+    # Filter out empty arrays for padding/concatenation
+    non_empty_obs = [obs for obs in all_obs_mx if obs.size > 0]
+    non_empty_acts = [act for act in all_acts_mx if act.size > 0]
+    non_empty_logprobs = [logprob for logprob in all_logprobs_mx if logprob.size > 0]
+
+    # Pad using native MLX operations
+    if non_empty_obs:
+        padded_obs = [mx.pad(obs, (0, max_obs_len - obs.shape[0]), constant_values=0)
+                      if obs.shape[0] < max_obs_len else obs[:max_obs_len]
+                      for obs in non_empty_obs]
+    else:
+        padded_obs = []
+
+    if non_empty_acts:
+        padded_acts = [mx.pad(act, (0, max_act_len - act.shape[0]), constant_values=0)
+                       if act.shape[0] < max_act_len else act[:max_act_len]
+                       for act in non_empty_acts]
+    else:
+        padded_acts = []
+
+    if non_empty_logprobs:
+        padded_logprobs = [mx.pad(logprob, (0, max_logprob_len - logprob.shape[0]), constant_values=0.0)
+                          if logprob.shape[0] < max_logprob_len else logprob[:max_logprob_len]
+                          for logprob in non_empty_logprobs]
+    else:
+        padded_logprobs = []
+
+    return {
+        'obs': mx.concatenate(padded_obs) if padded_obs else mx.array([], dtype=mx.int64),
+        'act': mx.concatenate(padded_acts) if padded_acts else mx.array([], dtype=mx.int64),
+        'logprob': mx.concatenate([lp.flatten() for lp in padded_logprobs]) if padded_logprobs else mx.array([], dtype=mx.float32),
         'rewards': mx.array(episode_rewards),
-        'episode_lengths': episode_lengths
+        'episode_lengths': episode_lengths,
     }
-    
+
+
+# Algorithm-specific data selection strategies
+def select_all_data(buffer):
+    """
+    GRPO data selector: Use all available data.
+
+    GRPO is on-policy but can benefit from using all collected episodes
+    since group-relative advantages normalize across the entire group.
+
+    Args:
+        buffer: Buffer containing episodes
+
+    Returns:
+        All episode data prepared for training
+    """
+    from textpolicy.buffer import Buffer
+    if not isinstance(buffer, Buffer):
+        raise TypeError(f"Expected Buffer, got {type(buffer)}")
+
+    episodes = buffer.episodes
+    if not episodes:
+        raise ValueError("Buffer is empty - no episodes to train on")
+
+    return _pack_episodes(episodes)
+
+
+def select_informative_data(buffer, min_variance: float = 0.01):
+    """
+    GRPO data selector with dynamic batch filtering (DAPO-style).
+
+    Filters out uninformative prompts where all completions have the same
+    outcome (all correct or all wrong), improving sample efficiency by
+    maintaining meaningful gradient signals.
+
+    This is the recommended selector for GRPO training when using multiple
+    completions per prompt, as it eliminates wasted compute on prompts
+    that provide no learning signal.
+
+    Args:
+        buffer: Buffer containing episodes (Episode objects or serialized dictionaries)
+        min_variance: Minimum reward variance to keep a prompt group.
+                     Prompts with variance below this are filtered out.
+
+    Returns:
+        Filtered episode data prepared for training, plus filtering stats.
+
+    Example:
+        >>> batch_data = select_informative_data(buffer, min_variance=0.01)
+        >>> # batch_data includes 'filter_stats' with filtering information
+
+    References:
+        DAPO: An Open-Source LLM Reinforcement Learning System at Scale
+        https://arxiv.org/abs/2503.14476
+    """
+    from textpolicy.buffer import Buffer
+    if not isinstance(buffer, Buffer):
+        raise TypeError(f"Expected Buffer, got {type(buffer)}")
+
+    episodes = buffer.episodes
+    if not episodes:
+        raise ValueError("Buffer is empty - no episodes to train on")
+
+    # Filter to keep only informative prompts
+    filtered_episodes, filter_stats = filter_informative_prompts(episodes, min_variance)
+
+    if not filtered_episodes:
+        raise ValueError(
+            f"All prompts filtered out (min_variance={min_variance}). "
+            f"Stats: {filter_stats}. Consider lowering min_variance or "
+            "ensuring diversity in completions."
+        )
+
+    # Pack filtered episodes using shared helper
+    batch_data = _pack_episodes(filtered_episodes)
+    batch_data['filter_stats'] = filter_stats
     return batch_data
 
 
 def select_recent_data(buffer, max_episodes: int = 100):
     """
     GRPO data selector: Use only recent episodes.
-    
+
     Alternative selector for GRPO that limits to recent episodes
     for faster training on large buffers.
-    
+
     Args:
         buffer: Buffer containing episodes (Episode objects or serialized dictionaries)
         max_episodes: Maximum number of recent episodes to use
-        
+
     Returns:
         Recent episode data prepared for training
     """
     from textpolicy.buffer import Buffer
     if not isinstance(buffer, Buffer):
         raise TypeError(f"Expected Buffer, got {type(buffer)}")
-    
+
     episodes = buffer.episodes
     if not episodes:
         raise ValueError("Buffer is empty - no episodes to train on")
-    
+
     # Select recent episodes
     recent_episodes = episodes[-max_episodes:] if len(episodes) > max_episodes else episodes
-    
-    # Process recent episodes
-    episode_rewards = []
-    episode_lengths = []
-    all_obs = []
-    all_acts = []
-    all_logprobs = []
-    
-    for episode in recent_episodes:
-        # Handle both Episode objects and serialized dictionaries
-        if hasattr(episode, 'rew'):
-            # Episode object with attributes
-            episode_reward = mx.sum(mx.array(episode.rew)).item()
-            episode_rewards.append(episode_reward)
-            episode_lengths.append(len(episode.obs))
-            
-            # For proper logprob extraction during training, we need the full context (prompt + response)
-            # This matches how the model was called during rollout generation
-            # Convert both obs and act to consistent Python list format before concatenation
-            obs_as_lists = []
-            for obs_item in episode.obs:
-                if hasattr(obs_item, 'tolist'):  # MLX array
-                    obs_as_lists.extend(obs_item.tolist())
-                elif isinstance(obs_item, list):  # Already Python list
-                    obs_as_lists.extend(obs_item)
-                else:  # Single item
-                    obs_as_lists.append(obs_item)
-            
-            act_as_lists = []
-            for act_item in episode.act:
-                if hasattr(act_item, 'tolist'):  # MLX array
-                    act_as_lists.extend(act_item.tolist())
-                elif isinstance(act_item, list):  # Already Python list
-                    act_as_lists.extend(act_item)
-                else:  # Single item
-                    act_as_lists.append(act_item)
-            
-            # Now concatenate the normalized lists
-            full_sequence = obs_as_lists + act_as_lists
-            all_obs.append(full_sequence)
-            
-            # Extract actions as consistent Python lists
-            episode_actions = []
-            for act_item in episode.act:
-                if hasattr(act_item, 'tolist'):  # MLX array
-                    episode_actions.extend(act_item.tolist())
-                elif isinstance(act_item, list):  # Already Python list
-                    episode_actions.extend(act_item)
-                else:  # Single item
-                    episode_actions.append(act_item)
-            all_acts.append(episode_actions)
-            
-            # Extract logprobs as consistent Python lists
-            episode_logprobs = []
-            if episode.logprob:
-                for logprob_item in episode.logprob:
-                    if hasattr(logprob_item, 'tolist'):  # MLX array
-                        episode_logprobs.extend(logprob_item.tolist())
-                    elif isinstance(logprob_item, list):  # Already Python list
-                        episode_logprobs.extend(logprob_item)
-                    else:  # Single item
-                        episode_logprobs.append(logprob_item)
-            all_logprobs.append(episode_logprobs)
-        else:
-            # Serialized dictionary from multiprocessing
-            episode_reward = mx.sum(episode['rew']).item()
-            episode_rewards.append(episode_reward)
-            episode_lengths.append(len(episode['obs']))
-            
-            # For proper logprob extraction during training, we need the full context (prompt + response)
-            # This matches how the model was called during rollout generation
-            # Convert both obs and act to consistent Python list format before concatenation
-            obs_as_lists = []
-            for obs_item in episode['obs']:
-                if hasattr(obs_item, 'tolist'):  # MLX array
-                    obs_as_lists.extend(obs_item.tolist())
-                elif isinstance(obs_item, list):  # Already Python list
-                    obs_as_lists.extend(obs_item)
-                else:  # Single item
-                    obs_as_lists.append(obs_item)
-            
-            act_as_lists = []
-            for act_item in episode['act']:
-                if hasattr(act_item, 'tolist'):  # MLX array
-                    act_as_lists.extend(act_item.tolist())
-                elif isinstance(act_item, list):  # Already Python list
-                    act_as_lists.extend(act_item)
-                else:  # Single item
-                    act_as_lists.append(act_item)
-            
-            # Now concatenate the normalized lists
-            full_sequence = obs_as_lists + act_as_lists
-            all_obs.append(full_sequence)
-            
-            # Extract actions as consistent Python lists
-            episode_actions = []
-            for act_item in episode['act']:
-                if hasattr(act_item, 'tolist'):  # MLX array
-                    episode_actions.extend(act_item.tolist())
-                elif isinstance(act_item, list):  # Already Python list
-                    episode_actions.extend(act_item)
-                else:  # Single item
-                    episode_actions.append(act_item)
-            all_acts.append(episode_actions)
-            
-            # Extract logprobs as consistent Python lists
-            episode_logprobs = []
-            if episode.get('logprob'):
-                for logprob_item in episode['logprob']:
-                    if hasattr(logprob_item, 'tolist'):  # MLX array
-                        episode_logprobs.extend(logprob_item.tolist())
-                    elif isinstance(logprob_item, list):  # Already Python list
-                        episode_logprobs.extend(logprob_item)
-                    else:  # Single item
-                        episode_logprobs.append(logprob_item)
-            all_logprobs.append(episode_logprobs)
-    
-    # Convert Python lists to MLX arrays before concatenation
-    # This is required because Episode objects store data as Python lists for memory efficiency
-    # For proper logprob extraction, we need uniform-length sequences, so we pad to the maximum length
-    
-    # Find maximum sequence length for padding
-    max_obs_len = max(len(obs) for obs in all_obs) if all_obs else 0
-    max_act_len = max(len(act) for act in all_acts) if all_acts else 0
-    max_logprob_len = max(len(logprob) for logprob in all_logprobs) if all_logprobs else 0
-    
-    # MLX-native padding and array operations for optimal Apple Silicon performance  
-    # Convert all sequences to MLX arrays and pad directly in MLX space
-    try:
-        # Convert all sequences to MLX arrays first (staying in unified memory)
-        all_obs_mx = [mx.array(obs, dtype=mx.int64) for obs in all_obs if obs]
-        all_acts_mx = [mx.array(act, dtype=mx.int64) for act in all_acts if act]  
-        all_logprobs_mx = [mx.array(logprob, dtype=mx.float32) for logprob in all_logprobs if logprob]
-        
-        # Pad using native MLX operations (more efficient for Apple Silicon)
-        if all_obs_mx:
-            padded_obs_mx = [mx.pad(obs, (0, max_obs_len - obs.shape[0]), constant_values=0) 
-                           if obs.shape[0] < max_obs_len else obs[:max_obs_len] 
-                           for obs in all_obs_mx]
-        else:
-            padded_obs_mx = []
-            
-        if all_acts_mx:
-            padded_acts_mx = [mx.pad(act, (0, max_act_len - act.shape[0]), constant_values=0) 
-                            if act.shape[0] < max_act_len else act[:max_act_len] 
-                            for act in all_acts_mx]
-        else:
-            padded_acts_mx = []
-            
-        if all_logprobs_mx:
-            padded_logprobs_mx = [mx.pad(logprob, (0, max_logprob_len - logprob.shape[0]), constant_values=0.0) 
-                                if logprob.shape[0] < max_logprob_len else logprob[:max_logprob_len] 
-                                for logprob in all_logprobs_mx]
-        else:
-            padded_logprobs_mx = []
-        
-        # Use padded MLX arrays directly (no intermediate conversion needed)
-        all_obs_mx = padded_obs_mx
-        all_acts_mx = padded_acts_mx
-        all_logprobs_mx = padded_logprobs_mx
-        
-    except Exception as e:
-        print(f"ERROR in MLX array conversion: {e}")
-        print(f"DEBUG: all_obs types: {[type(obs) for obs in all_obs[:3]]}")  # Show first 3 for brevity
-        print(f"DEBUG: all_logprobs types: {[type(logprob) for logprob in all_logprobs[:3]]}")
-        raise
-    
-    batch_data = {
-        'obs': mx.concatenate(all_obs_mx) if all_obs_mx else mx.array([]),  # Flat concatenated full sequences
-        'act': mx.concatenate(all_acts_mx) if all_acts_mx else mx.array([]),  # Flat concatenated response tokens
-        'logprob': mx.concatenate([logprob.flatten() for logprob in all_logprobs_mx]) if all_logprobs_mx else mx.array([]),  # Flat sequence for training
-        'rewards': mx.array(episode_rewards),
-        'episode_lengths': episode_lengths
-    }
-    
-    return batch_data
+
+    return _pack_episodes(recent_episodes)