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

textpolicy/rollout/worker.py

@@ -0,0 +1,194 @@
+# textpolicy/rollout/worker.py
+"""
+Worker process management and queue communication.
+"""
+
+import multiprocessing as mp
+import queue
+from typing import Callable, Any, Optional
+from .runner import RolloutRunner
+from .base import DEFAULT_WORKER_TIMEOUT, DEFAULT_MAX_STEPS
+from textpolicy.buffer import Buffer
+
+
+class RolloutWorker:
+    """
+    Manages a rollout runner in a separate process.
+    
+    Handles:
+    - Process lifecycle (start, stop, cleanup)
+    - Queue-based communication with trainer
+    - Data serialization for multiprocessing
+    - Async rollout collection
+    """
+    
+    def __init__(
+        self,
+        env_fn: Callable[[], Any],
+        policy_fn: Callable[[], Any],
+        strategy: Any,
+        max_steps: int = DEFAULT_MAX_STEPS,
+        send_queue: Optional[mp.Queue] = None,
+    ):
+        """
+        Initialize worker for separate process execution.
+
+        Args:
+            env_fn: Function that returns a fresh environment
+            policy_fn: Function that returns a fresh policy (for multiprocessing compatibility)
+            strategy: RolloutStrategy instance (e.g., PPOStrategy)
+            max_steps: Maximum steps per rollout collection
+            send_queue: Queue to send data back to trainer
+        """
+        self.env_fn = env_fn
+        self.policy_fn = policy_fn
+        self.strategy = strategy
+        self.max_steps = max_steps
+
+        # Communication queues
+        self.send_queue: mp.Queue = send_queue if send_queue is not None else mp.Queue()
+        self.control_queue: mp.Queue = mp.Queue()  # Commands: "collect", "exit"
+
+        # Process handle and status
+        self.process: Optional[mp.Process] = None
+        self.is_closed = False  # Track cleanup status
+
+    def run(self):
+        """
+        Target function for the worker process.
+        
+        Creates fresh environment and policy instances, then runs collection loop.
+        Communicates via queues to avoid shared memory issues.
+        """
+        # Create fresh instances in worker process (avoid pickle issues)
+        env = self.env_fn()
+        policy = self.policy_fn()
+        runner = RolloutRunner(env, policy, self.strategy, self.max_steps)
+
+        while True:
+            try:
+                # Wait for commands from trainer
+                msg = self.control_queue.get(timeout=DEFAULT_WORKER_TIMEOUT)
+                
+                if msg == "collect":
+                    # Run rollout collection
+                    buffer = runner.collect()
+                    
+                    # Serialize buffer to pure Python types for queue transmission
+                    episodes_data = [ep.to_dict() for ep in buffer.episodes]
+                    self.send_queue.put(episodes_data)
+                    
+                elif msg == "exit":
+                    break
+                    
+            except queue.Empty:
+                # Timeout - continue listening for commands
+                continue
+
+    def start(self):
+        """Launch the worker process."""
+        if self.process is not None:
+            raise RuntimeError("Worker process already started")
+            
+        self.process = mp.Process(target=self.run, daemon=True)
+        self.process.start()
+
+    def collect_async(self):
+        """Request a rollout without blocking."""
+        if self.process is None:
+            raise RuntimeError("Worker process not started")
+            
+        self.control_queue.put("collect")
+
+    def has_data(self) -> bool:
+        """Check if rollout result is ready."""
+        return not self.send_queue.empty()
+
+    def get_buffer(self) -> Buffer:
+        """
+        Retrieve the collected buffer.
+        
+        Call only after has_data() returns True.
+        
+        Returns:
+            Buffer instance with collected episodes
+        """
+        episodes_data = self.send_queue.get()
+        
+        # Reconstruct buffer from serialized data
+        buffer = Buffer(max_episodes=len(episodes_data))
+        for ep_dict in episodes_data:
+            buffer.add_episode_from_dict(ep_dict)
+            
+        return buffer
+
+    def close(self):
+        """Shut down the worker process gracefully."""
+        if self.process is None or self.is_closed:
+            return
+            
+        try:
+            # Step 1: Signal shutdown to worker process
+            self.control_queue.put("exit")
+            
+            # Step 2: Wait for graceful shutdown
+            self.process.join(timeout=3)
+            
+            # Step 3: Force termination if needed
+            if self.process.is_alive():
+                self.process.terminate()
+                self.process.join(timeout=1)
+            
+            # Step 4: Properly drain and close queues to prevent file descriptor errors
+            # This prevents the background feeder threads from trying to close already-closed FDs
+            self._cleanup_queues()
+                
+        except Exception:
+            # Ignore cleanup errors during testing - they're just noise
+            pass
+        finally:
+            # Step 5: Mark as closed
+            self.is_closed = True
+            self.process = None
+
+    def _cleanup_queues(self):
+        """
+        Properly cleanup multiprocessing queues to prevent file descriptor errors.
+        
+        The issue: Multiprocessing queues have background "feeder" threads that can
+        try to close file descriptors after the main process has already closed them.
+        
+        Solution: Explicitly drain and close queues before process termination.
+        """
+        try:
+            # Cancel the queue's join thread to prevent hanging on exit
+            # This tells the queue not to wait for the feeder thread
+            if hasattr(self.send_queue, 'cancel_join_thread'):
+                self.send_queue.cancel_join_thread()
+            if hasattr(self.control_queue, 'cancel_join_thread'):
+                self.control_queue.cancel_join_thread()
+                
+            # Drain any remaining items from queues to prevent deadlock
+            # Empty queues close more cleanly
+            try:
+                while not self.send_queue.empty():
+                    self.send_queue.get_nowait()
+            except queue.Empty:
+                pass  # Queue might be closed already
+                
+            try:
+                while not self.control_queue.empty():
+                    self.control_queue.get_nowait()
+            except queue.Empty:
+                pass  # Queue might be closed already
+                
+            # Now close the queues properly
+            if hasattr(self.send_queue, 'close'):
+                self.send_queue.close()
+            if hasattr(self.control_queue, 'close'):
+                self.control_queue.close()
+                
+        except Exception:
+            # During testing, ignore queue cleanup errors
+            # They're artifacts of rapid process termination
+            pass 

textpolicy/training/__init__.py

@@ -0,0 +1,14 @@
+# textpolicy/training/__init__.py
+"""
+Unified training infrastructure for all RL algorithms.
+"""
+
+from .trainer import Trainer
+from .rollout_manager import RolloutManager
+from .metrics import TrainingMetrics
+
+__all__ = [
+    "Trainer",
+    "RolloutManager", 
+    "TrainingMetrics"
+]

textpolicy/training/metrics.py

@@ -0,0 +1,242 @@
+# textpolicy/training/metrics.py
+"""
+Training metrics collection and analysis for all RL algorithms.
+"""
+
+from typing import Dict, Any, Optional, cast
+import mlx.core as mx # type: ignore
+from collections import defaultdict, deque
+
+
+class TrainingMetrics:
+    """
+    Lightweight metrics collector optimized for MLX training.
+    
+    Tracks algorithm-agnostic and algorithm-specific metrics
+    with minimal overhead during training.
+    """
+    
+    def __init__(self, history_length: int = 100):
+        """
+        Initialize metrics collector.
+        
+        Args:
+            history_length: Number of recent values to keep for rolling averages
+        """
+        self.history_length = history_length
+        self.metrics = defaultdict(lambda: deque(maxlen=history_length))
+        self.total_steps = 0
+        
+    def update(self, metrics_dict: Dict[str, float]):
+        """
+        Update metrics with new values.
+        
+        Args:
+            metrics_dict: Dictionary of metric_name -> value
+        """
+        for name, value in metrics_dict.items():
+            self.metrics[name].append(value)
+        
+        if 'step' in metrics_dict:
+            self.total_steps = metrics_dict['step']
+    
+    def get_latest(self, metric_name: str) -> Optional[float]:
+        """Get the most recent value for a metric."""
+        if metric_name in self.metrics and self.metrics[metric_name]:
+            return self.metrics[metric_name][-1]
+        return None
+    
+    def get_mean(self, metric_name: str, last_n: Optional[int] = None) -> Optional[float]:
+        """
+        Get mean of recent metric values.
+        
+        Args:
+            metric_name: Name of the metric
+            last_n: Number of recent values to average (None for all)
+            
+        Returns:
+            Mean value or None if metric doesn't exist
+        """
+        if metric_name not in self.metrics or not self.metrics[metric_name]:
+            return None
+        
+        values = list(self.metrics[metric_name])
+        if last_n is not None:
+            values = values[-last_n:]
+        
+        return sum(values) / len(values) if values else None
+    
+    def get_summary(self) -> Dict[str, Any]:
+        """
+        Get comprehensive metrics summary.
+        
+        Returns:
+            Dictionary with latest, mean, and other statistics
+        """
+        summary = {
+            'total_steps': self.total_steps,
+            'metrics': {}
+        }
+        
+        for metric_name, values in self.metrics.items():
+            if not values:
+                continue
+                
+            values_list = list(values)
+            summary['metrics'][metric_name] = {
+                'latest': values_list[-1],
+                'mean': sum(values_list) / len(values_list),
+                'min': min(values_list),
+                'max': max(values_list),
+                'count': len(values_list)
+            }
+        
+        return summary
+    
+    def reset(self):
+        """Reset all metrics."""
+        self.metrics.clear()
+        self.total_steps = 0
+    
+    def __len__(self) -> int:
+        """Return number of metrics being tracked."""
+        return len(self.metrics)
+
+
+class RolloutMetrics:
+    """
+    Metrics specific to rollout collection phase.
+    """
+    
+    def __init__(self):
+        self.episodes_collected = 0
+        self.total_reward = 0.0
+        self.episode_lengths = []
+        self.episode_rewards = []
+    
+    def add_episode(self, reward: float, length: int):
+        """Add metrics from a completed episode."""
+        self.episodes_collected += 1
+        self.total_reward += reward
+        self.episode_lengths.append(length)
+        self.episode_rewards.append(reward)
+    
+    def get_summary(self) -> Dict[str, float]:
+        """Get rollout metrics summary."""
+        if not self.episode_rewards:
+            return {
+                'episodes_collected': 0,
+                'mean_reward': 0.0,
+                'mean_length': 0.0,
+                'total_reward': 0.0
+            }
+        
+        return {
+            'episodes_collected': self.episodes_collected,
+            'mean_reward': sum(self.episode_rewards) / len(self.episode_rewards),
+            'mean_length': sum(self.episode_lengths) / len(self.episode_lengths),
+            'total_reward': self.total_reward,
+            'min_reward': min(self.episode_rewards),
+            'max_reward': max(self.episode_rewards),
+            'min_length': min(self.episode_lengths),
+            'max_length': max(self.episode_lengths)
+        }
+    
+    def reset(self):
+        """Reset rollout metrics."""
+        self.episodes_collected = 0
+        self.total_reward = 0.0
+        self.episode_lengths.clear()
+        self.episode_rewards.clear()
+
+
+def log_metrics(
+    metrics: Dict[str, float],
+    step: int,
+    logger: Optional[Any] = None,
+    prefix: str = ""
+):
+    """
+    Log metrics to console and optionally to external logger.
+    
+    Args:
+        metrics: Metrics dictionary
+        step: Training step number
+        logger: Optional external logger (wandb, tensorboard, etc.)
+        prefix: Prefix for metric names
+    """
+    # Console logging
+    metrics_str = ", ".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
+    print(f"Step {step} | {prefix}{metrics_str}")
+    
+    # External logger
+    if logger is not None and hasattr(logger, 'log'):
+        prefixed_metrics = {f"{prefix}{k}": v for k, v in metrics.items()}
+        logger.log(prefixed_metrics, step=step)
+
+
+def compute_explained_variance(predicted: mx.array, targets: mx.array) -> float:
+    """
+    Compute explained variance for value function evaluation.
+    
+    Explained variance = 1 - Var(targets - predicted) / Var(targets)
+    
+    Args:
+        predicted: Predicted values
+        targets: Target values
+        
+    Returns:
+        Explained variance (1.0 = perfect, 0.0 = no better than mean)
+    """
+    target_var = mx.var(targets)
+    if target_var == 0:
+        return 0.0
+    
+    residual_var = mx.var(targets - predicted)
+    explained_var = 1.0 - (residual_var / target_var)
+    
+    return cast(float, explained_var.item())  # MLX scalar array .item() returns Python float for float dtypes
+
+
+def compute_policy_metrics(
+    old_logprobs: mx.array,
+    new_logprobs: mx.array,
+    clip_ratio: float = 0.2
+) -> Dict[str, float]:
+    """
+    Compute standard policy optimization metrics.
+    
+    Args:
+        old_logprobs: Log probabilities from rollout
+        new_logprobs: Log probabilities from current policy
+        clip_ratio: Clipping ratio used in loss
+        
+    Returns:
+        Dictionary of policy metrics
+    """
+    # Importance ratio
+    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)
+    clip_fraction = mx.mean(clipped.astype(mx.float32))
+    
+    # KL divergence approximation
+    kl_div = mx.mean(old_logprobs - new_logprobs)
+    
+    # Policy change magnitude
+    ratio_mean = mx.mean(ratio)
+    ratio_std = mx.std(ratio)
+    
+    # Policy entropy (negative log probability mean)
+    entropy_mean = mx.mean(new_logprobs)
+    
+    return {
+        'policy/ratio_mean': cast(float, ratio_mean.item()),
+        'policy/ratio_std': cast(float, ratio_std.item()),
+        'policy/clip_fraction': cast(float, clip_fraction.item()),
+        'policy/kl_divergence': cast(float, kl_div.item()),
+        'policy/entropy': -cast(float, entropy_mean.item())  # Negative entropy: -E[log(p)]
+    }

textpolicy/training/rollout_manager.py

@@ -0,0 +1,78 @@
+# textpolicy/training/rollout_manager.py
+"""
+Lightweight rollout manager that integrates with existing rollout system.
+"""
+
+from typing import Callable, Any
+from textpolicy.rollout import RolloutCoordinator
+from textpolicy.buffer import Buffer
+from .metrics import RolloutMetrics
+
+
+class RolloutManager:
+    """
+    Simple manager that wraps the existing RolloutCoordinator 
+    with metrics tracking and convenient interface.
+    """
+    
+    def __init__(
+        self,
+        env_fn: Callable[[], Any],
+        policy_fn: Callable[[], Any],
+        algorithm: str = 'grpo',
+        num_workers: int = 0,
+        max_steps: int = 1000,
+        max_episodes: int = 100
+    ):
+        """
+        Initialize rollout manager.
+        
+        Args:
+            env_fn: Function that creates environment instances
+            policy_fn: Function that creates policy instances  
+            algorithm: Algorithm name ('grpo', 'ppo', etc.)
+            num_workers: Number of worker processes (0 = single-process)
+            max_steps: Maximum steps per rollout
+            max_episodes: Maximum episodes to buffer
+        """
+        self.coordinator = RolloutCoordinator(
+            env_fn=env_fn,
+            policy_fn=policy_fn,
+            algorithm=algorithm,
+            num_workers=num_workers,
+            max_steps=max_steps,
+            max_episodes=max_episodes
+        )
+        
+        self.metrics = RolloutMetrics()
+    
+    def collect(self) -> Buffer:
+        """
+        Collect rollout data and update metrics.
+        
+        Returns:
+            Buffer containing collected episodes
+        """
+        # Use existing rollout system
+        buffer = self.coordinator.collect()
+        
+        # Update metrics
+        for episode in buffer.episodes:
+            episode_data = episode.to_tensor_dict()
+            reward = episode_data['rew'].sum().item()
+            length = len(episode_data['obs'])
+            self.metrics.add_episode(reward, length)
+        
+        return buffer
+    
+    def get_metrics(self) -> dict:
+        """Get rollout collection metrics."""
+        return self.metrics.get_summary()
+    
+    def reset_metrics(self):
+        """Reset rollout metrics."""
+        self.metrics.reset()
+    
+    def close(self):
+        """Cleanup resources."""
+        self.coordinator.close()