### Install the package

Source: https://github.com/lucidrains/evolutionary-policy-optimization/blob/main/README.md

Install the library using pip.

```bash
$ pip install evolutionary-policy-optimization
```

--------------------------------

### Distributed Training Setup with Accelerate

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Shows how to set up distributed training using PyTorch distributed and HuggingFace Accelerate. Ensure Accelerate is installed and configured.

```python
import torch
from evolutionary_policy_optimization import create_agent, EPO, Env
from evolutionary_policy_optimization.distributed import (
    is_distributed,
    get_world_and_rank,
)

# Check distributed status
print(f"Distributed: {is_distributed()}")
world_size, rank = get_world_and_rank()
print(f"World size: {world_size}, Rank: {rank}")

# Create agent with Accelerate for automatic distributed setup
agent = create_agent(
    dim_state=512,
    num_latents=32,
    dim_latent=32,
    actor_num_actions=5,
    actor_dim=256,
    actor_mlp_depth=2,
    critic_dim=256,
    critic_mlp_depth=3,
    wrap_with_accelerate=True,  # Enable Accelerate integration
    accelerate_kwargs=dict(
        mixed_precision='fp16',  # Optional: mixed precision training
    ),
)

# EPO automatically handles distributed rollouts
epo = EPO(
    agent,
    episodes_per_latent=8,
    max_episode_length=100,
    fix_environ_across_latents=True,  # Sync env seeds across workers
)

env = Env((512,)) # Assuming Env is a wrapper for gymnasium.Env

# Training automatically parallelizes across workers
epo(agent, env, num_learning_cycles=50)
```

--------------------------------

### Initialize and use components

Source: https://github.com/lucidrains/evolutionary-policy-optimization/blob/main/README.md

Setup the latent gene pool, actor, and critic models, then perform a genetic algorithm update step.

```python
import torch

from evolutionary_policy_optimization import (
    LatentGenePool,
    Actor,
    Critic
)

latent_pool = LatentGenePool(
    num_latents = 128,
    dim_latent = 32,
)

state = torch.randn(1, 32)

actor = Actor(dim_state = 32, dim = 256, mlp_depth = 2, num_actions = 4, dim_latent = 32)
critic = Critic(dim_state = 32, dim = 256, mlp_depth = 3, dim_latent = 32)

latent = latent_pool(latent_id = 2)

actions = actor(state, latent)
value = critic(state, latent)

# interact with environment and receive rewards, termination etc

# derive a fitness score for each gene / latent

fitness = torch.randn(128)

latent_pool.genetic_algorithm_step(fitness) # update latent genes with genetic algorithm
```

--------------------------------

### Install test dependencies

Source: https://github.com/lucidrains/evolutionary-policy-optimization/blob/main/README.md

Install development dependencies for testing.

```bash
$ pip install '.[test]' # or `uv pip install '.[test]'`
```

--------------------------------

### LunarLander Training with Video Recording

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Complete example training an EPO agent on the LunarLander-v3 environment with video recording. This setup is useful for visualizing training progress. Ensure the 'recordings' directory exists or is created.

```python
import gymnasium as gym
from shutil import rmtree
from evolutionary_policy_optimization import EPO, GymnasiumEnvWrapper

# Setup LunarLander environment with video recording
gym_env = gym.make('LunarLander-v3', render_mode='rgb_array')

rmtree('./recordings', ignore_errors=True)
gym_env = gym.wrappers.RecordVideo(
    env=gym_env,
    video_folder='./recordings',
    name_prefix='lunar-video',
    episode_trigger=lambda eps_num: (eps_num % 250) == 0,
    disable_logger=True,
)

env = GymnasiumEnvWrapper(gym_env)
```

--------------------------------

### Instantiate and manage agent with create_agent

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Creates a complete EPO agent with configured Actor, Critic, and LatentGenePool. Includes methods for action sampling, value estimation, and checkpointing.

```python
import torch
from evolutionary_policy_optimization import create_agent

# Create a complete agent with all components configured
agent = create_agent(
    # State and action dimensions
    dim_state=512,
    actor_num_actions=5,

    # Latent gene pool configuration
    num_latents=16,           # Population size
    dim_latent=32,            # Gene dimension

    # Actor network architecture
    actor_dim=256,
    actor_mlp_depth=2,

    # Critic network architecture
    critic_dim=256,
    critic_mlp_depth=3,

    # Training hyperparameters
    use_critic_ema=True,           # Use exponential moving average for critic
    actor_lr=8e-4,                 # Actor learning rate
    critic_lr=8e-4,                # Critic learning rate
    latent_lr=1e-5,                # Latent gene learning rate (if not frozen)
    batch_size=32,                 # Batch size for PPO updates
    max_grad_norm=1.0,             # Gradient clipping

    # PPO loss configuration
    actor_loss_kwargs=dict(
        eps_clip=0.2,              # PPO clipping parameter
        entropy_weight=0.01,       # Entropy bonus for exploration
        norm_advantages=True       # Normalize advantages
    ),
    critic_loss_kwargs=dict(
        eps_clip=0.4               # Value function clipping
    ),

    # Gene pool evolution settings
    latent_gene_pool_kwargs=dict(
        frac_natural_selected=0.5,
        frac_tournaments=0.5,
        frozen_latents=True,       # Only evolve via GA, not gradients
    ),

    # Optional: HuggingFace Accelerate integration
    wrap_with_accelerate=True,
    accelerate_kwargs=dict(cpu=False),
)

# Get actions from the agent
state = torch.randn(1, 512)
latent_id = 0

# Get action logits
logits = agent.get_actor_actions(state, latent_id=latent_id)
print(f"Action logits: {logits.shape}")

# Sample action with temperature
action, log_prob = agent.get_actor_actions(
    state,
    latent_id=latent_id,
    sample=True,
    temperature=1.0
)
print(f"Sampled action: {action.item()}, log_prob: {log_prob.item():.4f}")

# Get value estimate
value = agent.get_critic_values(state, latent_id=latent_id)
print(f"Value estimate: {value.item():.4f}")

# Save and load agent checkpoints
agent.save('./checkpoint.pt', overwrite=True)
agent.load('./checkpoint.pt')
```

--------------------------------

### Create and Train Agent for LunarLander

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Demonstrates creating an agent optimized for the LunarLander environment, configuring EPO training, and saving/loading the agent.

```python
import gymnasium as gym
from evolutionary_policy_optimization import EPO

env = gym.make('LunarLander-v2')

# Create agent optimized for LunarLander
agent = env.to_epo_agent(
    num_latents=8,
    dim_latent=32,
    actor_dim=128,
    actor_mlp_depth=3,
    critic_dim=256,
    critic_mlp_depth=5,
    latent_gene_pool_kwargs=dict(
        frac_natural_selected=0.5,
        frac_tournaments=0.5,
    ),
    accelerate_kwargs=dict(cpu=False),
    actor_optim_kwargs=dict(cautious_factor=0.1),
    critic_optim_kwargs=dict(cautious_factor=0.1),
)

# Configure EPO training
epo = EPO(
    agent,
    episodes_per_latent=10,
    max_episode_length=250,
    action_sample_temperature=1.0,
)

# Train for 100 cycles
epo(agent, env, num_learning_cycles=100)

# Save final checkpoint
agent.save('./lunar_agent.pt', overwrite=True)

# Later: load and continue training or evaluate
agent.load('./lunar_agent.pt')
```

--------------------------------

### Perform end-to-end learning

Source: https://github.com/lucidrains/evolutionary-policy-optimization/blob/main/README.md

Configure the agent and EPO trainer to run learning cycles, including saving and loading the agent.

```python
import torch

from evolutionary_policy_optimization import (
    create_agent,
    EPO,
    Env
)

agent = create_agent(
    dim_state = 512,
    num_latents = 16,
    dim_latent = 32,
    actor_num_actions = 5,
    actor_dim = 256,
    actor_mlp_depth = 2,
    critic_dim = 256,
    critic_mlp_depth = 3,
    latent_gene_pool_kwargs = dict(
        frac_natural_selected = 0.5
    )
)

epo = EPO(
    agent,
    episodes_per_latent = 1,
    max_episode_length = 10,
    action_sample_temperature = 1.
)

env = Env((512,))

epo(agent, env, num_learning_cycles = 5)

# saving and loading

agent.save('./agent.pt', overwrite = True)
agent.load('./agent.pt')
```

--------------------------------

### GymnasiumEnvWrapper for Integration

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Provides seamless integration with OpenAI Gymnasium environments. Automatically extracts state and action dimensions for agent construction. Use this wrapper to easily adapt Gymnasium environments for EPO.

```python
import gymnasium as gym
from evolutionary_policy_optimization import EPO, GymnasiumEnvWrapper

# Create and wrap a Gymnasium environment
gym_env = gym.make('CartPole-v1')
env = GymnasiumEnvWrapper(gym_env)

# Automatically create agent matched to environment dimensions
agent = env.to_epo_agent(
    num_latents=8,
    dim_latent=32,
    actor_dim=128,
    actor_mlp_depth=3,
    critic_dim=256,
    critic_mlp_depth=5,
    latent_gene_pool_kwargs=dict(
        frac_natural_selected=0.5,
        frac_tournaments=0.5,
    ),
    accelerate_kwargs=dict(cpu=True),
)

# Get environment hyperparameters
env_hparams = env.to_agent_hparams()
print(f"State dimension: {env_hparams['dim_state']}")
print(f"Number of actions: {env_hparams['actor_num_actions']}")

# Create EPO trainer and train
epo = EPO(
    agent,
    episodes_per_latent=10,
    max_episode_length=500,
    action_sample_temperature=1.0,
)

epo(agent, env, num_learning_cycles=50)
```

--------------------------------

### EPO End-to-End Training Loop

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Orchestrates the complete training loop, handling environment rollouts, experience collection, and coordinated learning. Use this for a streamlined training process. Ensure agent and environment are created before initiating the loop.

```python
import torch
from evolutionary_policy_optimization import create_agent, EPO, Env

# Create agent
agent = create_agent(
    dim_state=512,
    num_latents=16,
    dim_latent=32,
    actor_num_actions=5,
    actor_dim=256,
    actor_mlp_depth=2,
    critic_dim=256,
    critic_mlp_depth=3,
    latent_gene_pool_kwargs=dict(
        frac_natural_selected=0.5,
        frac_tournaments=0.5,
    ),
    wrap_with_accelerate=False,
)

# Create EPO trainer
epo = EPO(
    agent=agent,
    episodes_per_latent=5,
    max_episode_length=100,
    action_sample_temperature=1.0,
    fix_environ_across_latents=True,  # Same env seeds for fair comparison
)

# Create mock environment (replace with your actual env)
env = Env(state_shape=(512,), can_terminate_after=10)

# Run complete training loop
epo(
    agent=agent,
    env=env,
    num_learning_cycles=10,  # Number of evolution + learning iterations
    seed=42,                 # Optional random seed for reproducibility
)

# Alternative: Manual training with fine-grained control
for cycle in range(10):
    # Gather experience from environment
    memories = epo.gather_experience_from(env)

    # Learn from collected experience (PPO updates + GA evolution)
    agent.learn_from(memories, epochs=2)

    print(f"Completed cycle {cycle + 1}")

# Save trained agent
agent.save('./trained_agent.pt', overwrite=True)
```

--------------------------------

### Initialize and use Critic for value estimation

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Configures a Critic network using HL-Gauss classification for value estimation. Supports both single-state and batch-mode inference.

```python
import torch
from evolutionary_policy_optimization import Critic

# Create a critic for value estimation
critic = Critic(
    dim_state=512,          # Environment observation dimension
    dim=256,                # Hidden layer dimension
    mlp_depth=4,            # Number of MLP layers
    dim_latent=32,          # Latent conditioning dimension
    use_regression=False,   # Use HL-Gauss classification (recommended)
    hl_gauss_loss_kwargs=dict(
        min_value=-10.0,    # Minimum expected return
        max_value=10.0,     # Maximum expected return
        num_bins=250        # Classification bins for value distribution
    )
)

# Estimate value for a state-latent pair
state = torch.randn(1, 512)
latent = torch.randn(1, 32)

value = critic(state, latent)
print(f"Estimated value: {value.item():.4f}")

# Get raw logits for training
value_logits = critic(state, latent, return_logits=True)
print(f"Value logits shape: {value_logits.shape}")  # torch.Size([1, 250])

# Batch value estimation
batch_states = torch.randn(32, 512)
batch_latents = torch.randn(32, 32)
batch_values = critic(batch_states, batch_latents)
print(f"Batch values shape: {batch_values.shape}")  # torch.Size([32])
```

--------------------------------

### Run tests

Source: https://github.com/lucidrains/evolutionary-policy-optimization/blob/main/README.md

Execute the test suite using pytest.

```bash
$ pytest tests/
```

--------------------------------

### Manage Latent Gene Populations with LatentGenePool

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Initializes a population of latent vectors and performs genetic algorithm or firefly optimization steps to evolve them based on fitness scores.

```python
import torch
from evolutionary_policy_optimization import LatentGenePool

# Create a gene pool with 128 latent genes of dimension 32
latent_pool = LatentGenePool(
    num_latents=128,              # Population size (number of genes)
    dim_latent=32,                # Dimension of each latent gene
    num_islands=4,                # Number of isolated sub-populations
    frac_natural_selected=0.25,   # Fraction of population to remove (least fit)
    frac_tournaments=0.25,        # Fraction of genes in each tournament
    frac_elitism=0.1,             # Top performers protected from mutation
    frac_migrate=0.1,             # Fraction that migrates between islands
    mutation_strength=1.0,        # Gaussian noise multiplier for mutations
    migrate_every=100,            # Steps between island migrations
    l2norm_latent=False,          # Normalize latents to hypersphere
    crossover_random=True,        # Random interpolation vs simple averaging
    fast_genetic_algorithm=False, # Use power-law mutation strength sampling
)

# Retrieve a specific latent gene by ID
latent = latent_pool(latent_id=5)
print(f"Latent shape: {latent.shape}")  # torch.Size([32])

# Retrieve multiple latents at once
batch_latent_ids = torch.tensor([0, 5, 10, 15])
latents = latent_pool(latent_id=batch_latent_ids)
print(f"Batch latents shape: {latents.shape}")  # torch.Size([4, 32])

# After evaluating fitness of all genes, run genetic algorithm step
fitness_scores = torch.randn(128)  # Higher = better fitness
did_update, new_genes = latent_pool.genetic_algorithm_step(fitness_scores)
print(f"GA updated population: {did_update}")

# Alternative: Firefly algorithm optimization step
latent_pool.firefly_step(
    fitness=fitness_scores,
    beta0=2.0,    # Exploitation factor
    gamma=1.0,    # Light intensity decay over distance
)

# Check population diversity via pairwise distances
distances = latent_pool.get_distance()
print(f"Diversity matrix shape: {distances.shape}")  # torch.Size([num_islands, latents_per_island, latents_per_island])
```

--------------------------------

### Maintain Plasticity with Shrink & Perturb

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Demonstrates using the `shrink_and_perturb_` function to maintain network plasticity during long training runs. This can also be configured automatically during agent creation.

```python
import torch
from evolutionary_policy_optimization import Actor
from evolutionary_policy_optimization.epo import shrink_and_perturb_

# Create and train an actor for a while...
actor = Actor(
    dim_state=512,
    num_actions=4,
    dim=256,
    mlp_depth=2,
    dim_latent=32,
)

# After many training steps, apply shrink and perturb to combat loss of plasticity
shrink_and_perturb_(
    actor,
    shrink_factor=0.5,    # Shrink weights toward zero (0 = no shrink, 1 = full reset)
    perturb_factor=0.01,  # Add this much noise to weights
)

# Can also be configured automatically in create_agent
from evolutionary_policy_optimization import create_agent

agent = create_agent(
    dim_state=512,
    num_latents=16,
    dim_latent=32,
    actor_num_actions=5,
    actor_dim=256,
    actor_mlp_depth=2,
    critic_dim=256,
    critic_mlp_depth=3,
    shrink_and_perturb_every=1000,  # Apply every 1000 learning steps
    shrink_and_perturb_kwargs=dict(
        shrink_factor=0.5,
        perturb_factor=0.01,
    ),
    wrap_with_accelerate=False,
)
```

--------------------------------

### Implement Latent-Conditioned Actor Networks

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Creates an Actor network that uses FiLM conditioning to process environment states alongside latent gene vectors for action selection.

```python
import torch
from evolutionary_policy_optimization import Actor

# Create an actor for a 512-dim state space with 4 discrete actions
actor = Actor(
    dim_state=512,      # Environment observation dimension
    num_actions=4,      # Number of discrete actions
    dim=256,            # Hidden layer dimension
    mlp_depth=3,        # Number of MLP layers
    dim_latent=32,      # Latent conditioning dimension
)

# Single state inference
state = torch.randn(1, 512)           # Batch of 1 state
latent = torch.randn(1, 32)           # Corresponding latent gene

action_logits = actor(state, latent)
print(f"Action logits shape: {action_logits.shape}")  # torch.Size([1, 4])

# Sample action from logits
probs = torch.softmax(action_logits, dim=-1)
action = torch.multinomial(probs, num_samples=1)
print(f"Selected action: {action.item()}")

# Batch inference with multiple states and latents
batch_states = torch.randn(32, 512)
batch_latents = torch.randn(32, 32)
batch_logits = actor(batch_states, batch_latents)
print(f"Batch action logits shape: {batch_logits.shape}")  # torch.Size([32, 4])
```

--------------------------------

### Import core Agent components

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Imports the primary classes required for manual Agent construction and state normalization.

```python
import torch
from evolutionary_policy_optimization import Actor, Critic, LatentGenePool
from evolutionary_policy_optimization.epo import Agent, StateNorm
```

--------------------------------

### Manual Agent Construction

Source: https://context7.com/lucidrains/evolutionary-policy-optimization/llms.txt

Construct an agent with full control over its components, including actor, critic, and latent gene pool. Ensure state normalization is applied correctly. The agent can be updated with fitness scores, and its device and latent gene count can be accessed.

```python
import torch
from evolutionary_policy_optimization import StateNorm, Actor, Critic, LatentGenePool, Agent

dim_state = 512
dim_latent = 32
num_latents = 64

state_norm = StateNorm(dim=dim_state)

actor = Actor(
    dim_state=dim_state,
    num_actions=4,
    dim=256,
    mlp_depth=2,
    dim_latent=dim_latent,
    state_norm=state_norm,
)

critic = Critic(
    dim_state=dim_state,
    dim=256,
    mlp_depth=3,
    dim_latent=dim_latent,
    state_norm=state_norm,
)

latent_pool = LatentGenePool(
    num_latents=num_latents,
    dim_latent=dim_latent,
    frac_natural_selected=0.25,
    frac_elitism=0.1,
)

agent = Agent(
    actor=actor,
    critic=critic,
    latent_gene_pool=latent_pool,
    state_norm=state_norm,
    use_critic_ema=True,
    critic_ema_beta=0.95,
    batch_size=64,
    diversity_aux_loss_weight=1e-3,  # Encourage latent diversity
    wrap_with_accelerate=False,
)

# Update gene pool with fitness scores
fitness = torch.randn(num_latents)
agent.update_latent_gene_pool_(fitness)

# Access device
print(f"Agent device: {agent.device}")

# Get number of latents
print(f"Number of latent genes: {agent.num_latents}")
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.