Eggroll Trainer¶

Eggroll Trainer is a library for Evolution Strategy (ES) trainers in PyTorch, featuring the EGGROLL algorithm—a novel ES method that provides a hundredfold increase in training speed over naïve evolution strategies.

Features¶

🔄 EGGROLL Algorithm - Low-rank perturbations for massive speedup
🧩 PyTorch Native - Works seamlessly with any PyTorch model
🎯 Flexible Architecture - Easy to subclass for custom ES algorithms
💾 Memory Efficient - Optimized for large population sizes
⚡ High Performance - Near-linear complexity in sequence length
🧪 Well-Tested - Comprehensive test suite included

What is EGGROLL?¶

EGGROLL (Evolution Guided General Optimization via Low-rank Learning) is a novel ES algorithm that uses low-rank perturbations instead of full-rank ones, dramatically reducing memory and computation requirements.

Key Innovation¶

For matrix parameters W ∈ R^(m×n), EGGROLL samples low-rank matrices: - A ∈ R^(m×r), B ∈ R^(n×r) where r << min(m,n) - Forms perturbations as A @ B.T

This reduces: - Memory: O(mn) → O(r(m+n)) - Computation: O(mn) → O(r(m+n))

Yet still achieves high-rank updates through population averaging!

Quick Start¶

import torch
import torch.nn as nn
from eggroll_trainer import EGGROLLTrainer

# Define a model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 1)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

# Define fitness function (higher is better)
def fitness_fn(model):
    # Your evaluation logic here
    return torch.randn(1).item()

# Create EGGROLL trainer
model = SimpleModel()
trainer = EGGROLLTrainer(
    model.parameters(),
    model=model,
    fitness_fn=fitness_fn,
    population_size=256,      # Large populations are efficient!
    learning_rate=0.01,
    sigma=0.1,
    rank=1,                   # Low-rank rank (1 is often sufficient)
    seed=42,
)

# Train
trainer.train(num_generations=100)

Installation¶

pipuv

pip install eggroll-trainer

uv add eggroll-trainer

For examples with plotting:

pip install "eggroll-trainer[examples]"
# or
uv add eggroll-trainer --extra examples

For development/contributing, see Contributing Guide.

3D Reinforcement Learning Examples¶

Train agents in 3D MuJoCo environments using EGGROLL:

Documentation¶

Getting Started - Installation and quick start guide
User Guide - Core concepts, trainers, and advanced usage
API Reference - Complete API documentation
Examples - Walkthroughs and tutorials
Research - Algorithm details and benchmarks

Key Concepts¶

Evolution Strategies¶

Evolution Strategies (ES) are a class of black-box optimization algorithms that: - Evaluate multiple perturbed versions of a model (population) - Use fitness scores to guide parameter updates - Don't require gradients—perfect for non-differentiable objectives

EGGROLL Advantages¶

Low-Rank Perturbations: For matrices, uses A @ B.T instead of full noise
Per-Layer Updates: Handles each parameter tensor independently
Fitness Normalization: Supports global or group-based normalization
Noise Reuse: Optional antithetic sampling for efficiency

References¶

License¶

This project is licensed under the MIT License - see the LICENSE file for details.