ezppo 0.0.0__py3-none-any.whl

ezppo/__init__.py

@@ -0,0 +1,29 @@
+"""
+EzPPO - A simple, unified PPO library for Discrete and Continuous environments.
+
+Usage:
+    from ezppo import DiscretePPOAgent, ContinuousPPOAgent
+
+# Discrete example (CartPole):
+    agent = DiscretePPOAgent(state_dim=4, action_dim=2)
+    action, log_prob = agent.act(state)
+
+# Continuous example (Pendulum):
+    agent = ContinuousPPOAgent(state_dim=3, action_dim=1)
+    action, log_prob, raw_action = agent.act(state)
+"""
+
+from .networks import DiscretePPO, ContinuousPPO, BasePPO
+from .normalizer import RunningNormalizer
+from .agents import BasePPOAgent, DiscretePPOAgent, ContinuousPPOAgent
+
+__version__ = "0.1.0"
+__all__ = [
+    "DiscretePPO",
+    "ContinuousPPO",
+    "BasePPO",
+    "RunningNormalizer",
+    "BasePPOAgent",
+    "DiscretePPOAgent",
+    "ContinuousPPOAgent",
+]

ezppo/agents.py

@@ -0,0 +1,475 @@
+"""
+Agent classes for Discrete and Continuous PPO.
+
+Both agents share a common base class with:
+- Trajectory collection (store/clear)
+- GAE computation
+- Observation normalization
+- Mini-batch training loop
+- Device management
+- Save/load checkpointing
+"""
+
+import torch
+import numpy as np
+import torch.optim as optim
+from torch.distributions import Categorical, Normal
+from torch.nn import MSELoss
+
+from .networks import DiscretePPO, ContinuousPPO
+from .normalizer import RunningNormalizer
+
+
+class BasePPOAgent:
+    """Shared base for both discrete and continuous agents."""
+
+    def __init__(self, network, lr=3e-4, gamma=0.99, gae_lambda=0.95,
+                 epsilon=0.2, critic_coef=0.5, entropy_coef=0.01,
+                 num_envs=1, trajectory_size=2048, device="cpu",
+                 normalize_observations=True, default_num_epochs=4,
+                 default_batch_size=64):
+        self.network = network
+        self.optimizer = optim.Adam(network.parameters(), lr=lr, eps=1e-5)
+
+        self.trajectory = []
+        self._num_envs = num_envs
+        self._trajectory_size = trajectory_size
+
+        # PPO hyperparameters
+        self.gamma = gamma
+        self.gae_lambda = gae_lambda
+        self.epsilon = epsilon
+        self.critic_coef = critic_coef
+        self.entropy_coef = entropy_coef
+        self.lr = lr
+        self.device = torch.device(device)
+
+        self.normalize_observations = normalize_observations
+        self.default_num_epochs = default_num_epochs
+        self.default_batch_size = default_batch_size
+
+        # Set up observation normalizer (optional)
+        if self.normalize_observations:
+            self.obs_normalizer = RunningNormalizer(network.input_dim)
+            self.obs_normalizer.to(self.device)
+        else:
+            self.obs_normalizer = None
+
+        self.network.to(self.device)
+
+    def store(self, packet: tuple) -> None:
+        """Store a transition in the trajectory."""
+        self.trajectory.append(packet)
+
+    def clear(self) -> None:
+        """Clear the trajectory buffer."""
+        self.trajectory = []
+
+    def _compute_gae(self, states, rewards, values, next_values, terminal_mask):
+        """
+        Compute Generalized Advantage Estimation (GAE).
+        
+        Args:
+            states: [T, B, obs_dim] or [T, obs_dim]
+            rewards: [T, B] or [T]
+            values: [T, B] or [T] - value of current states
+            next_values: [T, B] or [T] - value of next states
+            terminal_mask: [T, B] or [T] - 0 if terminal, 1 if not
+            
+        Returns:
+            advantages: [T, B] or [T]
+            critic_targets: [T, B] or [T] - advantages + values (for critic loss)
+        """
+        # Handle both single-env [T] and multi-env [T, B] cases
+        if rewards.ndim == 1:
+            T = rewards.shape[0]
+            advantages = torch.zeros(T, device=self.device)
+            last_advantage = 0
+            for t in reversed(range(T)):
+                td_error = rewards[t] + self.gamma * next_values[t] * terminal_mask[t] - values[t]
+                last_advantage = td_error + self.gamma * self.gae_lambda * last_advantage * terminal_mask[t]
+                advantages[t] = last_advantage
+        else:
+            T, B = rewards.shape
+            advantages = torch.zeros(T, B, device=self.device)
+            last_advantage = 0
+            for t in reversed(range(T)):
+                td_error = rewards[t] + self.gamma * next_values[t] * terminal_mask[t] - values[t]
+                last_advantage = td_error + self.gamma * self.gae_lambda * last_advantage * terminal_mask[t]
+                advantages[t] = last_advantage
+
+        critic_targets = advantages + values
+        return advantages, critic_targets
+
+    def _flatten_batch(self, *tensors):
+        """Flatten [T, B, ...] to [T*B, ...] for batched training."""
+        return [t.flatten(end_dim=1) for t in tensors]
+
+    def train(self, num_epochs=None, batch_size=None):
+        """
+        Train PPO on the collected trajectory.
+        
+        Args:
+            num_epochs: Number of PPO update epochs (uses default if None).
+            batch_size: Mini-batch size for training (uses default if None).
+        """
+        if not self.trajectory:
+            return
+
+        num_epochs = num_epochs if num_epochs is not None else self.default_num_epochs
+        batch_size = batch_size if batch_size is not None else self.default_batch_size
+
+        # Unpack trajectory
+        states, actions, rewards, old_log_probs, next_states, terminateds = zip(*self.trajectory)
+
+        # Convert to tensors
+        states = torch.from_numpy(np.stack(states)).float().to(self.device)
+        actions = torch.from_numpy(np.stack(actions)).float().to(self.device)
+        rewards = torch.from_numpy(np.stack(rewards)).float().to(self.device)
+        old_log_probs = torch.from_numpy(np.stack(old_log_probs)).float().to(self.device)
+        next_states = torch.from_numpy(np.stack(next_states)).float().to(self.device)
+        terminateds = torch.from_numpy(np.stack(terminateds)).float().to(self.device)
+
+        # Ensure consistent shapes
+        # For single env, stacking gives [T], [T, obs_dim] etc.
+        # We want [T, B] for scalar tensors and [T, B, obs_dim] for state tensors
+        if rewards.ndim == 1:
+            # Single environment case
+            rewards = rewards.unsqueeze(1)      # [T, 1]
+            old_log_probs = old_log_probs.unsqueeze(1)  # [T, 1]
+            terminateds = terminateds.unsqueeze(1)  # [T, 1]
+            actions = actions.unsqueeze(1)      # [T, 1]
+            if states.ndim == 2:
+                states = states.unsqueeze(1)    # [T, 1, obs_dim]
+            if next_states.ndim == 2:
+                next_states = next_states.unsqueeze(1)  # [T, 1, obs_dim]
+
+        terminal_mask = (1 - terminateds).float()
+
+        # Compute values
+        values = self.network.predict_value(states).squeeze(-1)
+        next_values = self.network.predict_value(next_states).squeeze(-1)
+
+        # Compute GAE
+        advantages, critic_targets = self._compute_gae(
+            states, rewards, values, next_values, terminal_mask
+        )
+
+        # Normalize advantages
+        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+
+        # Detach targets (they shouldn't be used for backprop)
+        old_log_probs = old_log_probs.detach()
+        advantages = advantages.detach()
+        critic_targets = critic_targets.detach()
+
+        # Flatten for batched training
+        states, actions, old_log_probs, critic_targets, advantages = \
+            self._flatten_batch(states, actions, old_log_probs, critic_targets, advantages)
+
+        total_samples = states.shape[0]
+        avg_policy_loss = 0
+        avg_critic_loss = 0
+        avg_entropy = 0
+        num_updates = 0
+
+        # Training loop
+        for _ in range(num_epochs):
+            indices = np.arange(total_samples)
+            np.random.shuffle(indices)
+
+            for start in range(0, total_samples, batch_size):
+                end = min(start + batch_size, total_samples)
+                batch_idx = indices[start:end]
+
+                batch_states = states[batch_idx]
+                batch_actions = actions[batch_idx]
+                batch_old_log_probs = old_log_probs[batch_idx]
+                batch_critic_targets = critic_targets[batch_idx]
+                batch_advantages = advantages[batch_idx]
+
+                # Get new log probs and entropy (implemented by subclass)
+                new_log_probs, entropy = self._get_log_prob_and_entropy(batch_states, batch_actions)
+
+                # Compute ratios
+                ratio = torch.exp(new_log_probs - batch_old_log_probs)
+                clipped_ratio = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon)
+
+                # Policy loss (clipped surrogate)
+                policy_loss = -torch.min(
+                    ratio * batch_advantages,
+                    clipped_ratio * batch_advantages
+                ).mean()
+
+                # Critic loss (MSE)
+                new_values = self.network.predict_value(batch_states).squeeze(-1)
+                critic_loss = MSELoss()(new_values, batch_critic_targets)
+
+                # Entropy bonus
+                entropy_loss = -entropy.mean()
+
+                # Total loss
+                total_loss = policy_loss + self.critic_coef * critic_loss + self.entropy_coef * entropy_loss
+
+                # Optimizer step
+                self.optimizer.zero_grad()
+                total_loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.network.parameters(), max_norm=0.5)
+                self.optimizer.step()
+
+                avg_policy_loss += policy_loss.item()
+                avg_critic_loss += critic_loss.item()
+                avg_entropy += entropy_loss.item()
+                num_updates += 1
+
+        # Update observation normalizer (if enabled)
+        if self.obs_normalizer is not None:
+            self.obs_normalizer.update(states.cpu().detach())
+
+        # Move back to CPU
+        self.network.to("cpu")
+        if self.obs_normalizer is not None:
+            self.obs_normalizer.to("cpu")
+
+        # Log
+        if num_updates > 0:
+            print(f"[{type(self).__name__}] Policy: {avg_policy_loss/num_updates:.4f} | "
+                  f"Critic: {avg_critic_loss/num_updates:.4f} | "
+                  f"Entropy: {avg_entropy/num_updates:.4f}")
+
+    def save(self, model_path: str, normalizer_path: str = None) -> None:
+        """Save checkpoint."""
+        torch.save(self.network.state_dict(), model_path)
+        if self.obs_normalizer is not None:
+            self.obs_normalizer.save(normalizer_path)
+
+    def load(self, model_path: str, normalizer_path: str | None = None) -> None:
+        """Load checkpoint.
+        
+        Args:
+            model_path: Path to the network state dict file.
+            normalizer_path: Path to the normalizer file (only used if
+                             normalize_observations=True; defaults to None).
+        """
+        self.network.load_state_dict(torch.load(model_path, weights_only=True))
+        if self.obs_normalizer is not None:
+            self.obs_normalizer.load(normalizer_path)
+        self.network.to(self.device)
+        if self.obs_normalizer is not None:
+            self.obs_normalizer.to(self.device)
+
+
+class DiscretePPOAgent(BasePPOAgent):
+    """
+    PPO agent for discrete action spaces.
+    
+    Usage:
+        agent = DiscretePPOAgent(state_dim=4, action_dim=2)
+        action, log_prob = agent.act(state)
+        agent.store((state, action, reward, log_prob, next_state, terminated))
+        agent.train()
+    """
+
+    def __init__(self, state_dim: int, action_dim: int,
+                 architecture: dict = None,
+                 lr: float = 3e-4, gamma: float = 0.99,
+                 gae_lambda: float = 0.95, epsilon: float = 0.2,
+                 critic_coef: float = 0.5, entropy_coef: float = 0.001,
+                 num_envs: int = 1, trajectory_size: int = 2048,
+                 device: str = "cpu", checkpoint: str = None,
+                 normalize_observations: bool = True,
+                 default_num_epochs: int = 4,
+                 default_batch_size: int = 64):
+        
+        # Default architecture
+        if architecture is None:
+            architecture = {
+                "actor": [state_dim, 64, 64, action_dim],
+                "critic": [state_dim, 64, 64, 1]
+            }
+
+        network = DiscretePPO(architecture)
+        super().__init__(network, lr, gamma, gae_lambda, epsilon,
+                         critic_coef, entropy_coef, num_envs,
+                         trajectory_size, device,
+                         normalize_observations, default_num_epochs,
+                         default_batch_size)
+
+        if checkpoint:
+            if self.obs_normalizer is not None:
+                self.load(checkpoint + ".pth", checkpoint + "_norm.npy")
+            else:
+                self.load(checkpoint + ".pth", None)
+
+    def act(self, state) -> tuple:
+        """
+        Select an action (stochastic).
+        
+        Args:
+            state: Observation (numpy array or tensor).
+            
+        Returns:
+            (action, log_prob)
+        """
+        state = torch.as_tensor(state, dtype=torch.float32)
+        if state.dim() == 1:
+            state = state.unsqueeze(0)
+        if self.normalize_observations:
+            state = self.obs_normalizer.normalize(state).to(self.device)
+        else:
+            state = state.to(self.device)
+
+        logits = self.network(state)
+        dist = Categorical(logits=logits)
+        action = dist.sample()
+        log_prob = dist.log_prob(action.squeeze(0))
+
+        return action.squeeze(0).cpu().detach(), log_prob.cpu().detach()
+
+    def act_deterministic(self, state) -> int:
+        """Select action greedily (no sampling)."""
+        state = torch.as_tensor(state, dtype=torch.float32)
+        if state.dim() == 1:
+            state = state.unsqueeze(0)
+        if self.normalize_observations:
+            state = self.obs_normalizer.normalize(state).to(self.device)
+        else:
+            state = state.to(self.device)
+
+        logits = self.network(state)
+        action = logits.argmax(dim=-1).squeeze(0)
+        return action.cpu().detach().item()
+
+    def _get_log_prob_and_entropy(self, states, actions):
+        """Compute log probability and entropy for a batch."""
+        logits = self.network(states)
+        dist = Categorical(logits=logits)
+        log_probs = dist.log_prob(actions)
+        entropies = dist.entropy().sum(-1)
+        return log_probs, entropies
+
+
+class ContinuousPPOAgent(BasePPOAgent):
+    """
+    PPO agent for continuous action spaces with tanh squashing.
+    
+    Uses reparameterization trick with tanh:
+    - Sample raw action from Gaussian
+    - Squash with tanh: action = tanh(raw)
+    - Correct log_prob with Jacobian: log_prob = raw_log_prob - log(1 - tanh^2 + eps)
+    
+    Usage:
+        agent = ContinuousPPOAgent(state_dim=3, action_dim=1)
+        action, log_prob, raw_action = agent.act(state)
+        agent.store((state, raw_action, reward, log_prob, next_state, terminated))
+        agent.train()
+    """
+
+    def __init__(self, state_dim: int, action_dim: int,
+                 architecture: dict = None,
+                 lr: float = 3e-4, gamma: float = 0.99,
+                 gae_lambda: float = 0.95, epsilon: float = 0.2,
+                 critic_coef: float = 0.5, entropy_coef: float = 0.01,
+                 num_envs: int = 1, trajectory_size: int = 2048,
+                 device: str = "cpu", checkpoint: str = None,
+                 normalize_observations: bool = True,
+                 default_num_epochs: int = 4,
+                 default_batch_size: int = 64):
+        
+        # Default architecture
+        if architecture is None:
+            architecture = {
+                "actor": [state_dim, 64, 64, action_dim],
+                "critic": [state_dim, 64, 64, 1]
+            }
+
+        network = ContinuousPPO(architecture)
+        super().__init__(network, lr, gamma, gae_lambda, epsilon,
+                         critic_coef, entropy_coef, num_envs,
+                         trajectory_size, device,
+                         normalize_observations, default_num_epochs,
+                         default_batch_size)
+
+        if checkpoint:
+            if self.obs_normalizer is not None:
+                self.load(checkpoint + ".pth", checkpoint + "_norm.npy")
+            else:
+                self.load(checkpoint + ".pth", None)
+
+    def act(self, state) -> tuple:
+        """
+        Select an action (stochastic, with tanh squashing).
+        
+        Args:
+            state: Observation (numpy array or tensor).
+            
+        Returns:
+            (action, log_prob, raw_action)
+            - action: Clamped to [-1, 1]
+            - log_prob: Corrected for tanh squashing
+            - raw_action: Unsquashed Gaussian sample
+        """
+        state = torch.as_tensor(state, dtype=torch.float32)
+        if state.dim() == 1:
+            state = state.unsqueeze(0)
+        if self.normalize_observations:
+            state = self.obs_normalizer.normalize(state).to(self.device)
+        else:
+            state = state.to(self.device)
+
+        mean, log_std = self.network(state)
+        log_std = torch.clamp(log_std, -5, 2)
+        std = log_std.exp()
+
+        dist = Normal(mean, std)
+        raw_action = dist.rsample()
+        action = torch.tanh(raw_action)
+
+        # Jacobian correction for tanh squashing
+        log_prob = (dist.log_prob(raw_action) - torch.log(
+            1 - action ** 2 + 1e-6
+        )).sum(-1)
+
+        return action.squeeze(0).cpu().detach(), \
+               log_prob.cpu().detach(), \
+               raw_action.squeeze(0).cpu().detach()
+
+    def act_deterministic(self, state) -> tuple:
+        """
+        Select action deterministically (mean of distribution).
+        Useful for evaluation.
+        """
+        state = torch.as_tensor(state, dtype=torch.float32)
+        if state.dim() == 1:
+            state = state.unsqueeze(0)
+        if self.normalize_observations:
+            state = self.obs_normalizer.normalize(state).to(self.device)
+        else:
+            state = state.to(self.device)
+
+        mean, log_std = self.network(state)
+        action = torch.tanh(mean)
+
+        log_prob = torch.zeros(mean.shape[0], device=self.device)
+        raw_action = torch.zeros(mean.shape[0], device=self.device)
+
+        return action.squeeze(0).cpu().detach(), \
+               log_prob.cpu().detach(), \
+               raw_action.cpu().detach()
+
+    def _get_log_prob_and_entropy(self, states, raw_actions):
+        """Compute log probability and entropy for a batch of raw actions."""
+        mean, log_std = self.network(states)
+        log_std = torch.clamp(log_std, -5, 2)
+        std = log_std.exp()
+
+        dist = Normal(mean, std)
+        actions = torch.tanh(raw_actions)
+
+        # Jacobian correction
+        log_probs = (dist.log_prob(raw_actions) - torch.log(
+            1 - actions ** 2 + 1e-6
+        )).sum(-1)
+
+        entropy = dist.entropy().sum(-1)
+        return log_probs, entropy

ezppo/networks.py

@@ -0,0 +1,145 @@
+"""
+PyTorch neural network modules for Discrete and Continuous PPO.
+
+Both networks share a common base class (BasePPO) that provides:
+- Configurable architecture via dict (actor, critic, optional shared backbone)
+- Orthogonal weight initialization (standard for RL)
+- Small output layer initialization (std=0.01) for stability
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class BasePPO(nn.Module):
+    """
+    Base PPO network with configurable architecture.
+    
+    Architecture dict format:
+        {
+            "actor": [input_dim, hidden1, hidden2, ..., output_dim],
+            "critic": [input_dim, hidden1, hidden2, ..., 1],
+            "shared": [input_dim, hidden1, ...]  # optional: shared backbone
+        }
+    
+    If "shared" is provided, both actor and critic share an initial backbone.
+    Otherwise, actor and critic are independent networks.
+    """
+
+    def __init__(self, architecture: dict, hidden_activation=nn.Tanh):
+        super().__init__()
+
+        self.actor_head = self._init_layer(architecture["actor"], hidden_activation)
+        self.critic_head = self._init_layer(architecture["critic"], hidden_activation)
+
+        self.backbone = self._init_layer(
+            architecture["shared"], hidden_activation
+        ) if "shared" in architecture else None
+
+        self.input_dim = architecture["shared"][0] if self.backbone else architecture["actor"][0]
+        self.output_dim = architecture["actor"][-1]
+
+        # Choose forward methods based on whether we have a shared backbone
+        if self.backbone:
+            self._act_forward = self._act_forward_with_backbone
+            self._crit_forward = self._crit_forward_with_backbone
+        else:
+            self._act_forward = self._act_forward_independent
+            self._crit_forward = self._crit_forward_independent
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Actor forward pass (returns logits for discrete, mean for continuous)."""
+        return self._act_forward(x)
+
+    def predict_value(self, x: torch.Tensor) -> torch.Tensor:
+        """Critic forward pass (returns state value estimate)."""
+        return self._crit_forward(x)
+
+    def _init_layer(self, arch: list[int], activation):
+        """Build a sequential network from an architecture list."""
+        seq = nn.Sequential()
+        for i, v in enumerate(arch):
+            if i == len(arch) - 1:
+                break  # Don't add output layer activation
+            input_n = v
+            output_n = arch[i + 1]
+
+            # Initialize weights with orthogonal init
+            if i + 1 == len(arch) - 1:
+                # Output layer: use small initialization for stability
+                linear = self._ortho_init(nn.Linear(input_n, output_n), std=0.01)
+            else:
+                # Hidden layers: use sqrt(2) orthogonal init (standard for RL)
+                linear = self._ortho_init(nn.Linear(input_n, output_n))
+
+            # Only add activation for hidden layers, not the output layer
+            if i + 1 != len(arch) - 1:
+                seq.extend([linear, activation()])
+            else:
+                seq.append(linear)
+        return seq
+
+    def _ortho_init(self, layer, std=np.sqrt(2), bias_const=0.0):
+        """Initialize a linear layer with orthogonal weights."""
+        nn.init.orthogonal_(layer.weight, std)
+        nn.init.constant_(layer.bias, bias_const)
+        return layer
+
+    # --- Forward method variants ---
+
+    def _act_forward_independent(self, x):
+        """Actor forward without shared backbone."""
+        return self.actor_head(x)
+
+    def _crit_forward_independent(self, x):
+        """Critic forward without shared backbone."""
+        return self.critic_head(x)
+
+    def _act_forward_with_backbone(self, x):
+        """Actor forward with shared backbone."""
+        return self.actor_head(self.backbone(x))
+
+    def _crit_forward_with_backbone(self, x):
+        """Critic forward with shared backbone."""
+        return self.critic_head(self.backbone(x))
+
+
+class DiscretePPO(BasePPO):
+    """
+    Discrete action PPO network.
+    
+    Output:
+        forward(x) -> action logits [batch, n_actions]
+        predict_value(x) -> value estimates [batch, 1]
+    """
+
+    def __init__(self, architecture: dict, hidden_activation=nn.Tanh):
+        super().__init__(architecture, hidden_activation)
+
+
+class ContinuousPPO(BasePPO):
+    """
+    Continuous action PPO network with learnable log_std.
+    
+    Output:
+        forward(x) -> (mean [batch, n_actions], log_std [n_actions])
+        predict_value(x) -> value estimates [batch, 1]
+    
+    The log_std parameter is shared across all batches and is clamped
+    during action sampling to prevent explosion.
+    """
+
+    def __init__(self, architecture: dict, hidden_activation=nn.Tanh, default_log=-1.0):
+        super().__init__(architecture, hidden_activation)
+
+        # Learnable log_std parameter (shared across all layers/batches)
+        self.log_std = nn.Parameter(
+            torch.ones(self.output_dim) * default_log,
+            requires_grad=True
+        )
+
+    def forward(self, x: torch.Tensor):
+        """Returns (mean, log_std). Mean has shape [batch, n_actions]."""
+        mean = super().forward(x)
+        return mean, self.log_std

ezppo/normalizer.py

@@ -0,0 +1,105 @@
+"""
+Running normalizer for observation preprocessing.
+Uses online algorithm to track running mean and variance.
+"""
+
+import torch
+import numpy as np
+
+
+class RunningNormalizer:
+    """
+    Online running normalizer for observations.
+    
+    Uses Welford's algorithm for numerically stable running statistics.
+    Clips normalized values to prevent extreme outliers.
+    """
+
+    def __init__(self, shape: int | list[int], epsilon: float = 1e-8, clip: float = 10.0):
+        """
+        Args:
+            shape: Dimension of observations (int or list of ints).
+                   Will be flattened to a 1D shape.
+            epsilon: Small constant for numerical stability.
+            clip: Maximum absolute value for normalized observations.
+        """
+        # Flatten shape if list or tuple
+        if isinstance(shape, (list, tuple)):
+            flat_shape = [int(np.prod(shape))]
+        else:
+            flat_shape = [int(shape)]
+
+        self.mean = torch.zeros(flat_shape)
+        self.var = torch.ones(flat_shape)
+        self.count = epsilon
+        self.epsilon = epsilon
+        self.clip = clip
+
+    def update(self, x: torch.Tensor) -> None:
+        """
+        Update running statistics with a batch of observations.
+        
+        Args:
+            x: Tensor of shape [batch, ...] or [batch, obs_dim]
+        """
+        batch_mean = torch.mean(x, dim=0)
+        batch_var = torch.var(x, dim=0, unbiased=False)
+        batch_count = x.shape[0]
+
+        total_count = self.count + batch_count
+        delta = batch_mean - self.mean
+
+        self.mean += delta * batch_count / total_count
+        self.var = (
+            (self.var * self.count + batch_var * batch_count +
+             delta ** 2 * self.count * batch_count / total_count)
+            / total_count
+        )
+        self.count = total_count
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Normalize observations.
+        
+        Args:
+            x: Tensor of observations.
+            
+        Returns:
+            Normalized tensor, clipped to [-clip, clip].
+        """
+        normed = (x - self.mean) / (torch.sqrt(self.var) + self.epsilon)
+        return torch.clip(normed, -self.clip, self.clip)
+
+    def denormalize(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Denormalize observations (inverse of normalize).
+        
+        Args:
+            x: Normalized tensor.
+            
+        Returns:
+            Original-scale tensor.
+        """
+        return (x * (torch.sqrt(self.var) + self.epsilon)) + self.mean
+
+    def save(self, path: str) -> None:
+        """Save statistics to a numpy file."""
+        np.save(path, [self.mean.cpu(), self.var.cpu()])
+
+    def load(self, path: str) -> None:
+        """Load statistics from a numpy file."""
+        data = np.load(path, allow_pickle=True)
+        self.mean = torch.as_tensor(data[0]).to(self.mean.device)
+        self.var = torch.as_tensor(data[1]).to(self.var.device)
+
+    def to(self, device: torch.device | str) -> "RunningNormalizer":
+        """Move statistics to device."""
+        self.mean = self.mean.to(device)
+        self.var = self.var.to(device)
+        return self
+
+    def reset(self) -> None:
+        """Reset statistics (useful for new training run)."""
+        self.mean = torch.zeros_like(self.mean)
+        self.var = torch.ones_like(self.var)
+        self.count = self.epsilon

ezppo-0.0.0.dist-info/METADATA

@@ -0,0 +1,149 @@
+Metadata-Version: 2.4
+Name: ezppo
+Version: 0.0.0
+Summary: A simple PPO library for Discrete and Continuous environments
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: gymnasium
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+
+# EzPPO
+
+A simple, unified PPO (Proximal Policy Optimization) library for PyTorch. Supports both **discrete** and **continuous** action spaces.
+
+## Features
+
+- **Unified API**: Same training loop for discrete and continuous
+- **Observation normalization**: Online running statistics
+- **GAE (Generalized Advantage Estimation)**: For stable advantage estimation
+- **Clipped surrogate objective**: Standard PPO clipping
+- **Mini-batch training**: With random shuffling
+- **Gradient clipping**: Prevents exploding gradients
+- **Checkpoint save/load**: Resume training easily
+- **Multi-environment support**: Parallel environments built-in
+- **GPU support**: Automatic device management
+
+## Quick Start
+
+### Discrete Actions (CartPole)
+
+```python
+from ezppo import DiscretePPOAgent
+import gymnasium as gym
+
+agent = DiscretePPOAgent(state_dim=4, action_dim=2)
+env = gym.make("CartPole-v1")
+
+state, _ = env.reset()
+action, log_prob = agent.act(state)  # Stochastic action
+# action = agent.act_deterministic(state)  # Greedy action (for eval)
+```
+
+### Continuous Actions (Pendulum)
+
+```python
+from ezppo import ContinuousPPOAgent
+import gymnasium as gym
+
+agent = ContinuousPPOAgent(state_dim=3, action_dim=1)
+env = gym.make("Pendulum-v1")
+
+state, _ = env.reset()
+action, log_prob, raw_action = agent.act(state)
+# action = agent.act_deterministic(state)[0]  # Mean action (for eval)
+```
+
+## Installation
+
+```bash
+pip install torch gymnasium numpy
+```
+
+## Architecture
+
+```
+ezppo/
+├── __init__.py          # Public API
+├── networks.py          # DiscretePPO, ContinuousPPO, BasePPO
+├── normalizer.py        # RunningNormalizer
+└── agents.py            # DiscretePPOAgent, ContinuousPPOAgent
+
+examples/
+├── discrete_cartpole.py
+└── continuous_pendulum.py
+```
+
+## Training Loop Template
+
+```python
+trajectory = []
+max_trajectory = 2048
+
+for episode in range(num_episodes):
+    state, _ = env.reset()
+    done = False
+
+    while not done:
+        action, log_prob = agent.act(state)
+        next_state, reward, terminated, truncated, _ = env.step(action.item())
+        done = terminated or truncated
+
+        trajectory.append((state, action.item(), reward, log_prob.item(), 
+                           next_state, terminated))
+        state = next_state
+
+    if len(trajectory) >= max_trajectory:
+        agent.train(num_epochs=4, batch_size=64)
+        trajectory.clear()
+```
+
+## API Reference
+
+### Networks
+
+| Class | Purpose |
+|-------|---------|
+| `BasePPO` | Base class with configurable architecture |
+| `DiscretePPO` | Discrete actions → logits output |
+| `ContinuousPPO` | Continuous actions → mean + learnable log_std |
+
+### Agents
+
+| Class | Purpose |
+|-------|---------|
+| `DiscretePPOAgent` | Discrete action space |
+| `ContinuousPPOAgent` | Continuous action space |
+
+### Key Methods
+
+**Agent.act(state)** → `(action, log_prob)` or `(action, log_prob, raw_action)`
+**Agent.act_deterministic(state)** → Deterministic action (for evaluation)
+**Agent.train(num_epochs=4, batch_size=64)** → Train on trajectory
+**Agent.store(packet)** → Add transition to trajectory
+**Agent.clear()** → Clear trajectory buffer
+**Agent.save(path)** → Save checkpoint
+**Agent.load(path)** → Load checkpoint
+
+## Configuration
+
+```python
+agent = DiscretePPOAgent(
+    state_dim=4, action_dim=2,
+    architecture={                    # Custom network architecture
+        "actor": [state_dim, 64, 64, action_dim],
+        "critic": [state_dim, 64, 64, 1]
+    },
+    lr=3e-4, gamma=0.99, 
+    gae_lambda=0.95, epsilon=0.2,
+    critic_coef=0.5, entropy_coef=0.01,
+    num_envs=1, trajectory_size=2048,
+    device="cpu"
+)
+```
+
+## License
+
+MIT

ezppo-0.0.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+ezppo/__init__.py,sha256=axJruk8hVTjLLzFJpSVqbjwfGhzbSASobOirgpKGJwY,786
+ezppo/agents.py,sha256=BnIqvczJpAUA0ySCUAZUUHQ7fuwl7V0FPvoaicc2AD0,18176
+ezppo/networks.py,sha256=HlzZMtFuHnLoctb3BnbVBdJvMPhT8a2Pgxr41ivQOE0,5261
+ezppo/normalizer.py,sha256=83pE-fbp0ts5aFSQFul6JF21tPnGNnuM9jYAqcay_k0,3398
+ezppo-0.0.0.dist-info/METADATA,sha256=5UAUJ5LMVdpT3xy1YyWOhmwaWsNWBxB8BeK4nqaZRPc,3980
+ezppo-0.0.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+ezppo-0.0.0.dist-info/top_level.txt,sha256=HwH4QV3kbktJE8cepSw2GIriNvXi66KwxfD6ubGeCgg,6
+ezppo-0.0.0.dist-info/RECORD,,

ezppo-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

ezppo-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ezppo