PufferLib - Apprentissage par renforcement haute performance

Vue d'ensemble

PufferLib est une bibliothèque d'apprentissage par renforcement haute performance conçue pour la simulation rapide d'environnements parallèles et l'entraînement. Elle atteint un entraînement à des millions d'étapes par seconde grâce à la vectorisation optimisée, au support natif multi-agent et à l'implémentation efficace de PPO (PuffeRL). La bibliothèque fournit la suite Ocean de 20+ environnements et une intégration transparente avec Gymnasium, PettingZoo et des frameworks RL spécialisés.

Quand utiliser cette skill

Utilisez cette skill pour :

Entraîner des agents RL avec PPO sur n'importe quel environnement (single-agent ou multi-agent)
Créer des environnements personnalisés en utilisant l'API PufferEnv
Optimiser les performances pour la simulation parallèle d'environnements (vectorisation)
Intégrer des environnements existants de Gymnasium, PettingZoo, Atari, Procgen, etc.
Développer des politiques avec CNN, LSTM ou architectures personnalisées
Mettre à l'échelle le RL à des millions d'étapes par seconde pour une expérimentation plus rapide
RL multi-agent avec support natif des environnements multi-agent

Capacités principales

1. Entraînement haute performance (PuffeRL)

PuffeRL est l'algorithme PPO+LSTM optimisé de PufferLib atteignant 1M-4M étapes/seconde.

Démarrage rapide de l'entraînement :

# Entraînement CLI
puffer train procgen-coinrun --train.device cuda --train.learning-rate 3e-4

# Entraînement distribué
torchrun --nproc_per_node=4 train.py

Boucle d'entraînement Python :

import pufferlib
from pufferlib import PuffeRL

# Créer un environnement vectorisé
env = pufferlib.make('procgen-coinrun', num_envs=256)

# Créer le trainer
trainer = PuffeRL(
    env=env,
    policy=my_policy,
    device='cuda',
    learning_rate=3e-4,
    batch_size=32768
)

# Boucle d'entraînement
for iteration in range(num_iterations):
    trainer.evaluate()  # Collecter les rollouts
    trainer.train()     # Entraîner sur le batch
    trainer.mean_and_log()  # Enregistrer les résultats

Pour des conseils d'entraînement complets, consultez references/training.md pour :

Flux de travail complet d'entraînement et options CLI
Ajustement des hyperparamètres avec Protein
Entraînement distribué multi-GPU/multi-nœud
Intégration de logger (Weights & Biases, Neptune)
Checkpointing et reprise d'entraînement
Conseils d'optimisation des performances
Modèles d'apprentissage par curriculum

2. Développement d'environnement (PufferEnv)

Créez des environnements personnalisés haute performance avec l'API PufferEnv.

Structure d'environnement basique :

import numpy as np
from pufferlib import PufferEnv

class MyEnvironment(PufferEnv):
    def __init__(self, buf=None):
        super().__init__(buf)

        # Définir les espaces
        self.observation_space = self.make_space((4,))
        self.action_space = self.make_discrete(4)

        self.reset()

    def reset(self):
        # Réinitialiser l'état et retourner l'observation initiale
        return np.zeros(4, dtype=np.float32)

    def step(self, action):
        # Exécuter l'action, calculer la récompense, vérifier done
        obs = self._get_observation()
        reward = self._compute_reward()
        done = self._is_done()
        info = {}

        return obs, reward, done, info

Utilisez le script template : scripts/env_template.py fournit des templates complets d'environnement single-agent et multi-agent avec exemples de :

Différents types d'espaces d'observation (vecteur, image, dict)
Variations d'espaces d'action (discret, continu, multi-discret)
Structure d'environnement multi-agent
Utilitaires de test

Pour le développement complet d'environnement, consultez references/environments.md pour :

Détails de l'API PufferEnv et modèles d'opération in-place
Définitions des espaces d'observation et d'action
Création d'environnements multi-agent
Suite Ocean (20+ environnements préconstruits)
Optimisation des performances (flux de travail Python à C)
Wrappers d'environnement et bonnes pratiques
Techniques de débogage et validation

3. Vectorisation et performance

Atteindre le débit maximum avec simulation parallèle optimisée.

Configuration de la vectorisation :

import pufferlib

# Vectorisation automatique
env = pufferlib.make('environment_name', num_envs=256, num_workers=8)

# Benchmarks de performance :
# - Envs en pur Python : 100k-500k SPS
# - Envs basés sur C : 100M+ SPS
# - Avec entraînement : 400k-4M SPS au total

Optimisations clés :

Buffers de mémoire partagée pour passage d'observation sans copie
Flags busy-wait au lieu de pipes/queues
Environnements excédentaires pour les retours asynchrones
Plusieurs environnements par worker

Pour l'optimisation de la vectorisation, consultez references/vectorization.md pour :

Architecture et caractéristiques de performance
Configuration des workers et taille des batches
Modes serial vs multiprocessing vs asynchrone
Modèles de mémoire partagée et sans copie
Vectorisation hiérarchique à grande échelle
Stratégies de vectorisation multi-agent
Profilage de performance et dépannage

4. Développement de politique

Construisez des politiques comme modules PyTorch standards avec utilitaires optionnels.

Structure de politique basique :

import torch.nn as nn
from pufferlib.pytorch import layer_init

class Policy(nn.Module):
    def __init__(self, observation_space, action_space):
        super().__init__()

        # Encoder
        self.encoder = nn.Sequential(
            layer_init(nn.Linear(obs_dim, 256)),
            nn.ReLU(),
            layer_init(nn.Linear(256, 256)),
            nn.ReLU()
        )

        # Têtes actor et critic
        self.actor = layer_init(nn.Linear(256, num_actions), std=0.01)
        self.critic = layer_init(nn.Linear(256, 1), std=1.0)

    def forward(self, observations):
        features = self.encoder(observations)
        return self.actor(features), self.critic(features)

Pour le développement complet de politique, consultez references/policies.md pour :

Politiques CNN pour observations d'images
Politiques récurrentes avec LSTM optimisé (inférence 3x plus rapide)
Politiques multi-entrées pour observations complexes
Politiques d'actions continues
Politiques multi-agent (paramètres partagés vs indépendants)
Architectures avancées (attention, résiduel)
Normalisation d'observation et gradient clipping
Débogage et test de politique

5. Intégration d'environnement

Intégrez facilement des environnements de frameworks RL populaires.

Intégration Gymnasium :

import gymnasium as gym
import pufferlib

# Wrapper d'environnement Gymnasium
gym_env = gym.make('CartPole-v1')
env = pufferlib.emulate(gym_env, num_envs=256)

# Ou utiliser make directement
env = pufferlib.make('gym-CartPole-v1', num_envs=256)

PettingZoo multi-agent :

# Environnement multi-agent
env = pufferlib.make('pettingzoo-knights-archers-zombies', num_envs=128)

Frameworks supportés :

Gymnasium / OpenAI Gym
PettingZoo (parallel et AEC)
Atari (ALE)
Procgen
NetHack / MiniHack
Minigrid
Neural MMO
Crafter
GPUDrive
MicroRTS
Griddly
Et bien d'autres...

Pour les détails d'intégration, consultez references/integration.md pour :

Exemples d'intégration complets pour chaque framework
Wrappers personnalisés (observation, récompense, frame stacking, action repeat)
Flattening et unflattening d'espaces
Enregistrement d'environnement
Modèles de compatibilité
Considérations de performance
Débogage d'intégration

Flux de travail de démarrage rapide

Pour entraîner des environnements existants

Choisir un environnement de la suite Ocean ou d'un framework compatible
Utiliser scripts/train_template.py comme point de départ
Configurer les hyperparamètres pour votre tâche
Exécuter l'entraînement avec CLI ou script Python
Monitorer avec Weights & Biases ou Neptune
Consulter references/training.md pour l'optimisation

Pour créer des environnements personnalisés

Commencer avec scripts/env_template.py
Définir les espaces d'observation et d'action
Implémenter les méthodes reset() et step()
Tester l'environnement localement
Vectoriser avec pufferlib.emulate() ou make()
Consulter references/environments.md pour les modèles avancés
Optimiser avec references/vectorization.md si nécessaire

Pour le développement de politique

Choisir une architecture basée sur les observations :
- Observations vecteur → politique MLP
- Observations image → politique CNN
- Tâches séquentielles → politique LSTM
- Observations complexes → politique multi-entrée
Utiliser layer_init pour l'initialisation appropriée des poids
Suivre les modèles dans references/policies.md
Tester avec l'environnement avant l'entraînement complet

Pour l'optimisation des performances

Profiler le débit actuel (étapes par seconde)
Vérifier la configuration de vectorisation (num_envs, num_workers)
Optimiser le code d'environnement (opérations in-place, vectorisation numpy)
Envisager une implémentation C pour les chemins critiques
Utiliser references/vectorization.md pour l'optimisation systématique

Ressources

scripts/

train_template.py - Template complet de script d'entraînement avec :

Création et configuration d'environnement
Initialisation de politique
Intégration de logger (WandB, Neptune)
Boucle d'entraînement avec checkpointing
Analyse d'arguments en ligne de commande
Configuration d'entraînement distribué multi-GPU

env_template.py - Templates d'implémentation d'environnement :

Exemple PufferEnv single-agent (grid world)
Exemple PufferEnv multi-agent (navigation coopérative)
Modèles multiples d'espaces d'observation/action
Utilitaires de test

references/

training.md - Guide complet d'entraînement :

Flux de travail d'entraînement et options CLI
Configuration des hyperparamètres
Entraînement distribué (multi-GPU, multi-nœud)
Monitoring et logging
Checkpointing
Ajustement des hyperparamètres avec Protein
Optimisation des performances
Modèles d'entraînement courants
Dépannage

environments.md - Guide de développement d'environnement :

API PufferEnv et caractéristiques
Espaces d'observation et d'action
Environnements multi-agent
Environnements de la suite Ocean
Flux de travail de développement d'environnement personnalisé
Chemin d'optimisation Python à C
Intégration d'environnement tiers
Wrappers et bonnes pratiques
Débogage

vectorization.md - Optimisation de la vectorisation :

Architecture et optimisations clés
Modes de vectorisation (serial, multiprocessing, asynchrone)
Configuration des workers et batches
Modèles de mémoire partagée et sans copie
Vectorisation avancée (hiérarchique, personnalisée)
Vectorisation multi-agent
Monitoring et profilage des performances
Dépannage et bonnes pratiques

policies.md - Guide d'architecture de politique :

Structure de politique basique
Politiques CNN pour images
Politiques LSTM avec optimisation
Politiques multi-entrées
Politiques d'actions continues
Politiques multi-agent
Architectures avancées (attention, résiduel)
Traitement d'observation et unflattening
Initialisation et normalisation
Débogage et test

integration.md - Guide d'intégration de framework :

Intégration Gymnasium
Intégration PettingZoo (parallel et AEC)
Environnements tiers (Procgen, NetHack, Minigrid, etc.)
Wrappers personnalisés (observation, récompense, frame stacking, etc.)
Conversion d'espaces et unflattening
Enregistrement d'environnement
Modèles de compatibilité
Considérations de performance
Débogage d'intégration

Conseils pour réussir

Commencez simplement : Commencez par les environnements Ocean ou l'intégration Gymnasium avant de créer des environnements personnalisés
Profilez tôt : Mesurez les étapes par seconde dès le début pour identifier les goulots d'étranglement
Utilisez les templates : scripts/train_template.py et scripts/env_template.py fournissent de solides points de départ
Lisez les références selon les besoins : Chaque fichier de référence est autonome et axé sur une capacité spécifique
Optimisez progressivement : Commencez avec Python, profilez, puis optimisez les chemins critiques avec C si nécessaire
Exploitez la vectorisation : La vectorisation de PufferLib est la clé pour atteindre un débit élevé
Monitorez l'entraînement : Utilisez WandB ou Neptune pour suivre les expériences et identifier les problèmes tôt
Testez les environnements : Validez la logique d'environnement avant d'augmenter l'échelle de l'entraînement
Vérifiez les environnements existants : La suite Ocean fournit 20+ environnements préconstruits
Utilisez l'initialisation appropriée : Utilisez toujours layer_init de pufferlib.pytorch pour les politiques

Cas d'usage courants

Entraînement sur des benchmarks standards

# Atari
env = pufferlib.make('atari-pong', num_envs=256)

# Procgen
env = pufferlib.make('procgen-coinrun', num_envs=256)

# Minigrid
env = pufferlib.make('minigrid-empty-8x8', num_envs=256)

Apprentissage multi-agent

# PettingZoo
env = pufferlib.make('pettingzoo-pistonball', num_envs=128)

# Politique partagée pour tous les agents
policy = create_policy(env.observation_space, env.action_space)
trainer = PuffeRL(env=env, policy=policy)

Développement de tâches personnalisées

# Créer un environnement personnalisé
class MyTask(PufferEnv):
    # ... implémenter l'environnement ...

# Vectoriser et entraîner
env = pufferlib.emulate(MyTask, num_envs=256)
trainer = PuffeRL(env=env, policy=my_policy)

Optimisation haute performance

# Maximiser le débit
env = pufferlib.make(
    'my-env',
    num_envs=1024,      # Grand batch
    num_workers=16,     # Nombreux workers
    envs_per_worker=64  # Optimiser par worker
)

Installation

uv pip install pufferlib

Documentation

Documentation officielle : https://puffer.ai/docs.html
GitHub : https://github.com/PufferAI/PufferLib
Discord : Support communautaire disponible