Mastering Pytorch Rl Nlp

Mastering PyTorch: Deep RL and NLP

Expert guidance for PyTorch 2.7+ development covering Deep Reinforcement Learning with TorchRL and NLP with HuggingFace Transformers.

Verified Library Versions (December 2025)

Deprecation Warnings

ALWAYS avoid these deprecated patterns:

Quick Reference

Device Setup (All Platforms)

import torch

device = (
    torch.device("cuda") if torch.cuda.is_available() else
    torch.device("mps") if torch.backends.mps.is_available() else
    torch.device("xpu") if hasattr(torch.backends, 'xpu') and torch.backends.xpu.is_available() else
    torch.device("cpu")
)
model = model.to(device)
model = torch.compile(model)  # Optimize for speed

TorchRL Quick Start (DQN)

from torchrl.envs import GymEnv
from torchrl.collectors import SyncDataCollector
from torchrl.data import ReplayBuffer, LazyTensorStorage
from torchrl.objectives import DQNLoss, HardUpdate

env = GymEnv("CartPole-v1", device=device)
collector = SyncDataCollector(env, policy, frames_per_batch=128, total_frames=100_000)
replay_buffer = ReplayBuffer(storage=LazyTensorStorage(10_000))
loss_module = DQNLoss(value_network=qnet, loss_function="smooth_l1", delay_value=True)

TorchRL Quick Start (PPO)

from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE

loss_fn = ClipPPOLoss(actor_network=actor, critic_network=critic, clip_epsilon=0.2, entropy_coef=0.01)
advantage_fn = GAE(value_network=critic, gamma=0.99, lmbda=0.95)

HuggingFace Quick Start

from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

# Simple inference
classifier = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
result = classifier("I love PyTorch!")

# Fine-tuning
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
    output_dir="./results", eval_strategy="epoch",  # NOT evaluation_strategy
    learning_rate=2e-5, per_device_train_batch_size=8, num_train_epochs=2, fp16=True
)
trainer = Trainer(model=model, args=training_args, train_dataset=train_ds, eval_dataset=val_ds)
trainer.train()

Detailed Guides

For comprehensive coverage, load the appropriate guide:

Common Patterns

Gymnasium Environment (Modern API)

import gymnasium as gym

env = gym.make("CartPole-v1")
obs, info = env.reset()

while True:
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)  # 5 values!
    done = terminated or truncated
    if done:
        break

Training Loop with AMP

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for batch in dataloader:
    optimizer.zero_grad()
    with autocast():
        loss = compute_loss(batch)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

PEFT/LoRA Fine-tuning

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"])
model = get_peft_model(model, lora_config)
# Now fine-tune with much fewer parameters

Installation

# Create virtual environment
python -m venv .venv && source .venv/bin/activate

# PyTorch with CUDA 12.4
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# RL libraries
pip install torchrl gymnasium pettingzoo

# NLP libraries
pip install transformers datasets peft accelerate

# Experiment tracking
pip install tensorboard wandb

Apple Silicon (M2/M3/M4) Support

PyTorch MPS backend enables GPU acceleration on Apple Silicon Macs.

MPS Setup

import torch
import os

# Enable MPS fallback for unsupported operations
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
model = model.to(device)

MPS Limitations (Dec 2025)

• Not all operations supported - Some fall back to CPU
• No distributed training - Single GPU only
• No float64 support - Use float32
• SDPA can be unstable - Use eager attention if issues occur
• torch.compile limited - Test thoroughly, may need to disable
• ~3x slower than RTX 4090 - But 80% lower energy consumption

Best for Apple Silicon

• Prototyping and development
• Light to medium training workloads
• Local inference
• Learning and experimentation

Key Concepts

Reinforcement Learning

• Agent: Learns from environment interactions
• Environment: Provides states, rewards (use Gymnasium, not gym)
• Policy: Maps states to actions (actor in PPO)
• Value Function: Estimates future rewards (critic in PPO)
• Experience Replay: Stores transitions for stable learning (DQN)
• Target Network: Slowly-updated copy for stable Q-targets (DQN)

NLP/Transformers

• Tokenizer: Converts text to token IDs
• Encoder: Processes input (BERT-style)
• Decoder: Generates output (GPT-style)
• Fine-tuning: Adapts pre-trained model to specific task
• PEFT/LoRA: Parameter-efficient fine-tuning (fewer trainable params)

Optimization

• torch.compile: JIT compilation for faster execution
• Mixed Precision (AMP): fp16/bf16 for speed and memory
• Quantization: Reduce model size (int8)
• DDP: Distributed training across GPUs
• torchrun: Launch distributed training

Mastering PyTorch: Deep RL and NLP

This repository contains a Codex skill that provides expert guidance for PyTorch development, focused on deep reinforcement learning with TorchRL and NLP transformers with HuggingFace.

Contents

• SKILL.md defines the skill metadata and core guidance.
• references/ includes deeper topic references linked from the skill.
• AGENTS.md outlines contributor guidelines for updates.

Usage

Open SKILL.md to review or update the skill content. When editing, keep examples concise and aligned with the verified library versions listed in the document.

Installing with Skilz (Universal Installer)

The recommended way to install this skill across different AI coding agents is using the skilz universal installer.

Install Skilz

pip install skilz

This skill supports Agent Skill Standard which means it supports 14 plus coding agents including Claude Code, OpenAI Codex, Cursor and Gemini.

Git URL Options

You can use either -g or --git with HTTPS or SSH URLs:

# HTTPS URL
skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill>

# SSH URL
skilz install --git git@github.com:SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill.git

Claude Code

Install to user home (available in all projects):

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill>

Install to current project only:

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill> --project

OpenCode

Install for OpenCode:

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill> --agent opencode

Project-level install:

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill> --project --agent opencode

Gemini

Project-level install for Gemini:

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill> --agent gemini

OpenAI Codex

Install for OpenAI Codex:

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill> --agent codex

Project-level install:

skilz install -g <https://github.com/SpillwaveSolutions/mastering-pytorch-rl-nlp-agentic-skill> --project --agent codex

Install from Skillzwave Marketplace

# Claude to user home dir ~/.claude/skills
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp

# Claude skill in project folder ./claude/skills
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp --project

# OpenCode install to user home dir ~/.config/opencode/skills
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp --agent opencode

# OpenCode project level
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp --agent opencode --project

# OpenAI Codex install to user home dir ~/.codex/skills
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp

# OpenAI Codex project level ./.codex/skills
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp --agent opencode --project

# Gemini CLI (project level) -- only works with project level
skilz install SpillwaveSolutions_mastering-pytorch-rl-nlp-agentic-skill/mastering-pytorch-rl-nlp --agent gemini

See this site skill listing to see how to install this exact skill to 14+ different coding agents.

Other Supported Agents

Skilz supports 14+ coding agents including Claude Code, OpenAI Codex, OpenCode, Cursor, Gemini CLI, GitHub Copilot CLI, Windsurf, Qwen Code, Aidr, and more.

For the full list of supported platforms, visit SkillzWave.ai/platforms or see the skilz-cli GitHub repository

<a href="https://skillzwave.ai/">Largest Agentic Marketplace for AI Agent Skills</a> and <a href="https://spillwave.com/">SpillWave: Leaders in AI Agent Development.</a>

Contributing

Please read AGENTS.md before making changes.

Advanced Topics Reference

Multi-modal learning, CLIP, RLHF, and responsible AI practices.

Multi-Modal Learning

CLIP (Contrastive Language-Image Pre-training)

from transformers import CLIPProcessor, CLIPModel
import torch
from PIL import Image

# Load model
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Process image and text
image = Image.open("image.jpg")
texts = ["a photo of a cat", "a photo of a dog", "a photo of a bird"]

inputs = processor(
    text=texts,
    images=image,
    return_tensors="pt",
    padding=True
)

# Get similarity scores
outputs = model(**inputs)
logits_per_image = outputs.logits_per_image
probs = logits_per_image.softmax(dim=1)
print(f"Probabilities: {probs}")

Image-Text Matching

from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer

# Image captioning
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

image = Image.open("image.jpg")
pixel_values = processor(images=image, return_tensors="pt").pixel_values

output_ids = model.generate(pixel_values, max_length=16)
caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)

RLHF (Reinforcement Learning from Human Feedback)

Overview

RLHF trains language models to align with human preferences: 1. Supervised Fine-Tuning (SFT) - Fine-tune on demonstration data 2. Reward Modeling - Train model to predict human preferences 3. RL Fine-Tuning - Use PPO to optimize for reward model

Using TRL Library

pip install trl

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer

# Load model with value head
model = AutoModelForCausalLMWithValueHead.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Configure PPO
ppo_config = PPOConfig(
    batch_size=16,
    learning_rate=1e-5,
    mini_batch_size=4,
    ppo_epochs=4
)

ppo_trainer = PPOTrainer(
    config=ppo_config,
    model=model,
    tokenizer=tokenizer
)

Reward Model Training

from transformers import AutoModelForSequenceClassification, Trainer

# Load reward model
reward_model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=1  # Scalar reward
)

# Train on preference data
# Dataset format: (prompt, chosen_response, rejected_response)

Ethics and Responsible AI

Bias Detection

# Check for bias in embeddings
def measure_bias(model, tokenizer, word_pairs):
    """Measure association bias between word pairs."""
    results = {}
    for target, attributes in word_pairs.items():
        target_emb = get_embedding(model, tokenizer, target)
        attr_embs = [get_embedding(model, tokenizer, a) for a in attributes]
        results[target] = compute_association(target_emb, attr_embs)
    return results

# Example word pairs for gender bias
word_pairs = {
    "programmer": ["he", "she"],
    "nurse": ["he", "she"],
}

Fairness Metrics

from sklearn.metrics import confusion_matrix

def demographic_parity(predictions, sensitive_attribute):
    """Check if positive prediction rates are equal across groups."""
    groups = {}
    for pred, attr in zip(predictions, sensitive_attribute):
        if attr not in groups:
            groups[attr] = []
        groups[attr].append(pred)
    
    rates = {g: sum(p) / len(p) for g, p in groups.items()}
    return rates

def equalized_odds(predictions, labels, sensitive_attribute):
    """Check if TPR and FPR are equal across groups."""
    # Implementation for equalized odds metric
    pass

Model Cards

Always document your models with:

• Intended use cases
• Training data sources
• Performance metrics by demographic
• Known limitations and biases
• Ethical considerations

# Model Card: MyClassifier

## Model Details
- Model type: BERT-based classifier
- Training data: [Dataset name and description]
- Training procedure: Fine-tuned for 3 epochs

## Intended Use
- Primary: Sentiment analysis for product reviews
- Out of scope: Medical or legal advice

## Limitations
- May not perform well on non-English text
- Biased toward certain demographics

## Ethical Considerations
- Tested for fairness across demographic groups
- Should not be used for automated decisions without human review

Robustness Testing

# Adversarial examples
def perturb_text(text, perturbation_type="typo"):
    """Generate adversarial perturbations."""
    if perturbation_type == "typo":
        # Add random typos
        pass
    elif perturbation_type == "synonym":
        # Replace with synonyms
        pass
    return perturbed_text

# Test model robustness
def test_robustness(model, test_cases):
    original_preds = model(test_cases)
    perturbed = [perturb_text(t) for t in test_cases]
    perturbed_preds = model(perturbed)
    consistency = sum(o == p for o, p in zip(original_preds, perturbed_preds))
    return consistency / len(test_cases)

Advanced Transformer Architectures

Efficient Attention

# Scaled Dot-Product Attention (SDPA) - PyTorch native
import torch.nn.functional as F

# Uses Flash Attention when available
output = F.scaled_dot_product_attention(
    query, key, value,
    attn_mask=mask,
    dropout_p=0.1 if training else 0.0,
    is_causal=True  # For autoregressive models
)

FlexAttention (PyTorch 2.5+)

from torch.nn.attention.flex_attention import flex_attention

# Custom attention patterns
def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

output = flex_attention(query, key, value, score_mod=causal_mask)

Multi-Agent RL

PettingZoo Environments

from pettingzoo.mpe import simple_spread_v3
from torchrl.envs import PettingZooEnv

# Create multi-agent environment
env = PettingZooEnv(
    env=simple_spread_v3.parallel_env(),
    categorical_actions=True
)

# Each agent gets observations and takes actions
td = env.reset()
for agent in env.agents:
    print(f"{agent}: {td[agent]['observation'].shape}")

Independent Learners

# Simple approach: train each agent independently
agents = {name: PPOAgent(obs_dim, act_dim) for name in env.agents}

for episode in range(num_episodes):
    td = env.reset()
    while not td["done"].all():
        actions = {}
        for name, agent in agents.items():
            actions[name] = agent.act(td[name]["observation"])
        td = env.step(actions)
        
        for name, agent in agents.items():
            agent.update(td[name])

Best Practices Summary

1. Multi-Modal: Use pre-trained models like CLIP, fine-tune for specific tasks 2. RLHF: Start with SFT, then reward modeling, finally PPO tuning 3. Ethics: Always test for bias, document limitations, use model cards 4. Robustness: Test with adversarial examples, edge cases 5. Efficiency: Use native SDPA, FlexAttention for custom patterns 6. Multi-Agent: Start with independent learners, then centralized training

NLP and Transformers Reference

Comprehensive guide to NLP with HuggingFace Transformers 4.56+, BERT, GPT, fine-tuning, and PEFT/LoRA.

HuggingFace Transformers Setup

Installation

pip install transformers datasets peft accelerate

Core Imports

from transformers import (
    AutoTokenizer,
    AutoModel,
    AutoModelForSequenceClassification,
    AutoModelForCausalLM,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    pipeline,
    DataCollatorWithPadding
)
from datasets import load_dataset

Pipeline API (Quick Inference)

The fastest way to use pre-trained models:

from transformers import pipeline

# Sentiment Analysis
classifier = pipeline("sentiment-analysis")
result = classifier("I love PyTorch!")
# [{'label': 'POSITIVE', 'score': 0.9998}]

# Specific model
classifier = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english"
)

# Text Generation
generator = pipeline("text-generation", model="gpt2")
result = generator("PyTorch is", max_length=50)

# Question Answering
qa = pipeline("question-answering")
result = qa(
    question="What is PyTorch?",
    context="PyTorch is a machine learning framework based on Torch."
)

# Named Entity Recognition
ner = pipeline("ner", grouped_entities=True)
result = ner("Hugging Face is based in New York City.")

# Zero-shot Classification
classifier = pipeline("zero-shot-classification")
result = classifier(
    "This is a great movie!",
    candidate_labels=["positive", "negative", "neutral"]
)

# Summarization
summarizer = pipeline("summarization")
result = summarizer(long_text, max_length=100, min_length=30)

# Translation
translator = pipeline("translation_en_to_fr")
result = translator("Hello, how are you?")

Pipeline with GPU

import torch

device = 0 if torch.cuda.is_available() else -1  # -1 for CPU
classifier = pipeline("sentiment-analysis", device=device)

# Or with specific device
classifier = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english",
    device="cuda:0"
)

Tokenizers

Basic Tokenization

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# Single text
tokens = tokenizer("Hello, world!")
print(tokens)
# {'input_ids': [...], 'attention_mask': [...]}

# Batch processing
texts = ["Hello, world!", "PyTorch is great!"]
tokens = tokenizer(
    texts,
    padding=True,
    truncation=True,
    max_length=512,
    return_tensors="pt"  # Return PyTorch tensors
)

# Decode back to text
decoded = tokenizer.decode(tokens["input_ids"][0])

Advanced Tokenization

# For sequence classification (single text)
tokens = tokenizer(
    text,
    padding="max_length",
    truncation=True,
    max_length=128,
    return_tensors="pt"
)

# For sequence pair tasks (e.g., NLI, QA)
tokens = tokenizer(
    text_a,
    text_b,
    padding=True,
    truncation=True,
    return_tensors="pt"
)

# Token-level tasks (NER)
tokens = tokenizer(
    text,
    is_split_into_words=True,  # Pre-tokenized input
    return_offsets_mapping=True
)

Models

Loading Models

from transformers import AutoModel, AutoModelForSequenceClassification
import torch

# Base model (embeddings only)
model = AutoModel.from_pretrained("bert-base-uncased")

# Classification head
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2
)

# Causal LM (text generation)
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("gpt2")

# Move to device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

Model Inference

model.eval()
with torch.no_grad():
    inputs = tokenizer(text, return_tensors="pt").to(device)
    outputs = model(**inputs)

    # Classification
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=-1)

    # Get hidden states
    hidden_states = outputs.last_hidden_state  # [batch, seq_len, hidden_dim]

Fine-Tuning with Trainer

Complete Fine-Tuning Example

from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    DataCollatorWithPadding
)
from datasets import load_dataset
import torch

# Load dataset
dataset = load_dataset("imdb")

# Load tokenizer and model
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=2
)

# Tokenize dataset
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512
    )

tokenized_datasets = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"]
)

# Data collator for dynamic padding
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# Training arguments - NOTE: use "epoch" for when to run validation
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",  # When to validate
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    warmup_ratio=0.1,
    logging_dir="./logs",
    logging_steps=100,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    fp16=True,  # Mixed precision
    push_to_hub=False
)

# Compute metrics
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

def compute_metrics(pred_tuple):
    logits, labels = pred_tuple
    predictions = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy_score(labels, predictions),
        "f1": f1_score(labels, predictions, average="weighted")
    }

# Initialize Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics
)

# Train
trainer.train()

# Validate
results = trainer.evaluate()
print(results)

# Save model
trainer.save_model("./final_model")
tokenizer.save_pretrained("./final_model")

PEFT and LoRA

Parameter-Efficient Fine-Tuning dramatically reduces trainable parameters.

Installation

pip install peft

LoRA Fine-Tuning

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForSequenceClassification

# Load base model
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2
)

# Configure LoRA
lora_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=8,  # Rank
    lora_alpha=16,  # Scaling factor
    lora_dropout=0.1,
    target_modules=["query", "value"],  # Which layers to adapt
    bias="none"
)

# Apply LoRA
model = get_peft_model(model, lora_config)

# Check trainable parameters
model.print_trainable_parameters()
# trainable params: 294,912 || all params: 109,778,690 || trainable%: 0.27%

# Train as normal with Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=validation_dataset,
    tokenizer=tokenizer
)
trainer.train()

# Save LoRA weights only
model.save_pretrained("./lora_model")

# Load LoRA model
from peft import PeftModel
base_model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
model = PeftModel.from_pretrained(base_model, "./lora_model")

QLoRA (Quantized LoRA)

For larger models with limited memory:

from transformers import BitsAndBytesConfig
import torch

# Quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True
)

# Load quantized model
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto"
)

# Apply LoRA on quantized model
model = get_peft_model(model, lora_config)

Token Classification (NER)

from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification
)
from datasets import load_dataset

# Load dataset
dataset = load_dataset("conll2003")

# Load model
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(
    model_name,
    num_labels=len(dataset["train"].features["ner_tags"].feature.names)
)

# Tokenize with label alignment
def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True
    )

    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        label_ids = []
        previous_word_idx = None
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)  # Ignore in loss
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(-100)  # Subword tokens
            previous_word_idx = word_idx
        labels.append(label_ids)

    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_dataset = dataset.map(tokenize_and_align_labels, batched=True)

# Data collator
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

# Train
trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="./ner_model",
        evaluation_strategy="epoch",
        learning_rate=2e-5,
        num_train_epochs=3
    ),
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer
)
trainer.train()

Text Generation

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Generate text
input_text = "PyTorch is"
inputs = tokenizer(input_text, return_tensors="pt")

# Generation parameters
outputs = model.generate(
    **inputs,
    max_new_tokens=50,
    num_return_sequences=1,
    temperature=0.7,
    top_p=0.9,
    top_k=50,
    do_sample=True,
    pad_token_id=tokenizer.eos_token_id
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

Popular Models

Classification

• bert-base-uncased - General purpose
• distilbert-base-uncased - Faster, smaller
• microsoft/deberta-v3-base - Best accuracy
• roberta-base - Robust BERT

Generation

• gpt2, gpt2-medium, gpt2-large
• meta-llama/Llama-2-7b-hf
• mistralai/Mistral-7B-v0.1

Embeddings

• sentence-transformers/all-MiniLM-L6-v2 - Fast
• sentence-transformers/all-mpnet-base-v2 - Quality

Multilingual

• bert-base-multilingual-cased
• xlm-roberta-base

Best Practices

Memory Optimization

# Gradient checkpointing
model.gradient_checkpointing_enable()

# Mixed precision
training_args = TrainingArguments(
    fp16=True,  # NVIDIA GPUs
    bf16=True,  # Ampere+ GPUs, better for transformers
)

# Gradient accumulation
training_args = TrainingArguments(
    gradient_accumulation_steps=4,  # Effective batch = batch_size * 4
)

Efficient Inference

# Batch inference
texts = ["text1", "text2", "text3"]
inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs.to(device))

# Use torch.compile
model = torch.compile(model)

Data Loading

from datasets import load_dataset

# Stream large datasets
dataset = load_dataset("wikipedia", streaming=True)

# Efficient preprocessing
dataset = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=4,  # Parallel processing
    remove_columns=["text"]
)

Common Issues

Out of Memory

• Reduce batch size
• Enable gradient checkpointing
• Use gradient accumulation
• Use PEFT/LoRA
• Use mixed precision (fp16/bf16)

Slow Training

• Use DataLoader with num_workers > 0
• Use pin_memory=True
• Ensure GPU utilization is high
• Use torch.compile()

Poor Results

• Verify data quality
• Try different learning rates (1e-5 to 5e-5)
• Increase training epochs
• Use learning rate warmup
• Check for label imbalance

Optimization and Deployment Reference

Comprehensive guide to torch.compile, quantization, distributed training, and production deployment.

torch.compile (PyTorch 2.0+)

Basic Usage

import torch

model = MyModel()
model = model.to(device)

# Compile for faster execution
model = torch.compile(model)

# Now use normally
output = model(input)

Compilation Modes

# Default - balanced compile time and performance
model = torch.compile(model, mode="default")

# Reduce overhead - faster compile, good for smaller models
model = torch.compile(model, mode="reduce-overhead")

# Max autotune - slowest compile, fastest runtime
model = torch.compile(model, mode="max-autotune")

Platform Support

Mixed Precision Training (AMP)

Automatic Mixed Precision

from torch.cuda.amp import autocast, GradScaler
import torch

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()

    # Forward pass in mixed precision
    with autocast():
        output = model(batch)
        loss = criterion(output, targets)

    # Scaled backward pass
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

BFloat16 (Ampere+ GPUs)

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    output = model(input)

Quantization

Dynamic Quantization

import torch

quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},
    dtype=torch.qint8
)

HuggingFace Quantization

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto"
)

Distributed Training

DistributedDataParallel (DDP)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

Launch with torchrun

torchrun --nproc_per_node=4 train.py

ONNX Export

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17,
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
)

TorchServe Deployment

pip install torchserve torch-model-archiver

torch-model-archiver \
    --model-name my_model \
    --version 1.0 \
    --serialized-file model.pth \
    --handler handler.py

torchserve --start --model-store model_store --models my_model=my_model.mar

Inference Optimization Checklist

1. Use torch.compile - 10-30% speedup 2. Enable mixed precision - Faster and less memory 3. Batch inputs - Better GPU utilization 4. Use torch.no_grad() - No gradient computation 5. Quantize if possible - Smaller, faster 6. Use model.train(False)

model.train(False)
model = model.to(device)
model = torch.compile(model)

@torch.no_grad()
def predict(inputs):
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        return model(inputs.to(device))

PyTorch Fundamentals Reference

Comprehensive guide to PyTorch 2.7+ fundamentals including tensors, autograd, nn.Module, and device management.

Device Management (Multi-Platform)

Universal Device Detection

import torch

def get_device():
    """Get the best available device for computation."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available():
        return torch.device("mps")
    elif hasattr(torch.backends, 'xpu') and torch.backends.xpu.is_available():
        return torch.device("xpu")
    return torch.device("cpu")

device = get_device()
print(f"Using device: {device}")

Apple Silicon (M2/M3/M4) MPS Backend

Status (December 2025): MPS backend is officially supported in PyTorch 1.12+.

# Check MPS availability
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print("MPS device available")
else:
    print("MPS not available, using CPU")

Known MPS Limitations:

• Not all operations implemented - some fall back to CPU
• No distributed training support
• No float64 or fp16 tensor core support
• SDPA (Scaled Dot-Product Attention) can be unstable

MPS Best Practices:

import os

# Enable MPS fallback for unsupported operations
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

# For transformer models with attention issues
os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"  # Prevents memory spikes

Performance Notes:

• ~3x slower than RTX 4090 for large training
• 80% lower energy consumption
• Best for: prototyping, inference, light-medium training

CUDA Device Management

import torch

# Check CUDA availability
if torch.cuda.is_available():
    print(f"CUDA available: {torch.cuda.get_device_name(0)}")
    print(f"CUDA version: {torch.version.cuda}")

# Multi-GPU selection
device = torch.device("cuda:0")  # First GPU
device = torch.device("cuda:1")  # Second GPU

# Memory management
torch.cuda.empty_cache()
print(f"Memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

Tensors

Creating Tensors

import torch

# From Python data
x = torch.tensor([1, 2, 3], dtype=torch.float32)
x = torch.tensor([[1, 2], [3, 4]], device=device)

# Common creation functions
zeros = torch.zeros(3, 4, device=device)
ones = torch.ones(3, 4, dtype=torch.float16)
rand = torch.rand(3, 4)  # Uniform [0, 1)
randn = torch.randn(3, 4)  # Normal distribution
arange = torch.arange(0, 10, 2)  # [0, 2, 4, 6, 8]
linspace = torch.linspace(0, 1, 5)  # 5 evenly spaced

# From NumPy (shares memory!)
import numpy as np
np_array = np.array([1, 2, 3])
tensor = torch.from_numpy(np_array)

Tensor Operations

# Basic math
a = torch.tensor([1, 2, 3], dtype=torch.float32)
b = torch.tensor([4, 5, 6], dtype=torch.float32)

c = a + b  # Element-wise addition
c = a * b  # Element-wise multiplication
c = a @ b  # Dot product (1D) or matrix multiply (2D+)
c = torch.matmul(a.unsqueeze(0), b.unsqueeze(1))  # Explicit matmul

# In-place operations (append _)
a.add_(1)  # Modifies a in place
a.zero_()  # Zeros in place

# Reduction operations
x = torch.randn(3, 4)
x.sum()  # Total sum
x.sum(dim=0)  # Sum along rows (result: shape [4])
x.sum(dim=1)  # Sum along columns (result: shape [3])
x.mean(), x.std(), x.var()
x.max(), x.min(), x.argmax(), x.argmin()

Reshaping and Indexing

x = torch.randn(2, 3, 4)

# Reshaping
x.view(6, 4)  # Must be contiguous
x.reshape(6, 4)  # Works on non-contiguous
x.flatten()  # To 1D
x.squeeze()  # Remove dims of size 1
x.unsqueeze(0)  # Add dim at position 0

# Indexing
x[0]  # First element along dim 0
x[:, 1]  # All rows, second column
x[..., -1]  # Last element along last dim
x[x > 0]  # Boolean indexing

# Advanced indexing
indices = torch.tensor([0, 2])
x.index_select(dim=1, index=indices)

Autograd

Gradient Computation

import torch

# Enable gradient tracking
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = x ** 2 + 3 * x
z = y.sum()

# Compute gradients
z.backward()
print(x.grad)  # dz/dx = 2x + 3 = [7.0, 9.0]

# Disable gradient tracking
with torch.no_grad():
    y = x * 2  # No gradient computation

# Detach from computation graph
y = x.detach()  # Same data, no gradient tracking

Gradient Control

# Zero gradients (important before each backward!)
optimizer.zero_grad()
# Or manually:
x.grad.zero_()

# Gradient clipping (prevents exploding gradients)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

# Gradient accumulation
for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

nn.Module

Building Models

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

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

# Using nn.Sequential
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(256, 10)
)

Model Operations

# Move to device
model = model.to(device)

# Training vs evaluation mode
model.train()  # Enable dropout, batch norm training mode
model.eval()   # Disable dropout, use running stats for batch norm

# Parameter access
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

# Freeze parameters
for param in model.fc1.parameters():
    param.requires_grad = False

Saving and Loading

# Save model state dict (recommended)
torch.save(model.state_dict(), "model.pth")

# Load model state dict
model = SimpleNet(input_dim, hidden_dim, output_dim)
model.load_state_dict(torch.load("model.pth", map_location=device))

# Save entire model (less portable)
torch.save(model, "model_full.pth")
model = torch.load("model_full.pth")

# Save checkpoint with optimizer state
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, "checkpoint.pth")

Common Layers

Linear and Activation

import torch.nn as nn

# Linear layers
nn.Linear(in_features, out_features, bias=True)

# Activations
nn.ReLU()
nn.LeakyReLU(negative_slope=0.01)
nn.GELU()  # Popular in transformers
nn.SiLU()  # Swish activation
nn.Tanh()
nn.Sigmoid()
nn.Softmax(dim=-1)

Normalization

# Batch normalization (for CNNs, MLPs)
nn.BatchNorm1d(num_features)
nn.BatchNorm2d(num_features)

# Layer normalization (for transformers, RNNs)
nn.LayerNorm(normalized_shape)

# Group normalization
nn.GroupNorm(num_groups, num_channels)

Regularization

# Dropout
nn.Dropout(p=0.5)
nn.Dropout2d(p=0.5)  # For conv layers

# Weight decay (L2 regularization) - applied via optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

Data Loading

Dataset and DataLoader

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

dataset = CustomDataset(X, y)
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,  # Faster GPU transfer
    drop_last=True    # Drop incomplete batches
)

for batch_data, batch_labels in dataloader:
    batch_data = batch_data.to(device)
    batch_labels = batch_labels.to(device)
    # Train...

Training Loop Pattern

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

# Setup
model = SimpleNet(input_dim, hidden_dim, output_dim).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)

# Training loop
for epoch in range(num_epochs):
    model.train()
    total_loss = 0

    for batch_data, batch_labels in train_loader:
        batch_data = batch_data.to(device)
        batch_labels = batch_labels.to(device)

        optimizer.zero_grad()
        outputs = model(batch_data)
        loss = criterion(outputs, batch_labels)
        loss.backward()

        # Optional: gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        total_loss += loss.item()

    scheduler.step()

    # Validation
    model.eval()
    with torch.no_grad():
        val_loss = 0
        correct = 0
        for batch_data, batch_labels in val_loader:
            batch_data = batch_data.to(device)
            batch_labels = batch_labels.to(device)
            outputs = model(batch_data)
            val_loss += criterion(outputs, batch_labels).item()
            correct += (outputs.argmax(dim=1) == batch_labels).sum().item()

        accuracy = correct / len(val_loader.dataset)

    print(f"Epoch {epoch}: Train Loss={total_loss:.4f}, Val Acc={accuracy:.4f}")

torch.compile (PyTorch 2.0+)

import torch

# Basic compilation (recommended for most cases)
model = torch.compile(model)

# Compilation modes
model = torch.compile(model, mode="default")      # Balanced
model = torch.compile(model, mode="reduce-overhead")  # Fast, less memory
model = torch.compile(model, mode="max-autotune")     # Slowest compile, fastest run

# Note: torch.compile has limited MPS support as of Dec 2025
# For Apple Silicon, test thoroughly or use without compile

Best Practices Summary

1. Always use device-agnostic code - check CUDA, MPS, XPU availability 2. Use `torch.no_grad()` for inference - saves memory and computation 3. Zero gradients before backward - prevent gradient accumulation bugs 4. Use `model.train()` and `model.eval()` - affects dropout and batch norm 5. Pin memory for GPU training - faster data transfer 6. Save state_dict, not full model - more portable 7. Use AMP for faster training - see optimization guide

Reinforcement Learning with TorchRL Reference

Comprehensive guide to Deep Reinforcement Learning using TorchRL 0.10.x, Gymnasium 1.0+, DQN, and PPO.

Critical: Use Gymnasium, NOT OpenAI Gym

OpenAI Gym is DEPRECATED. Always use Gymnasium:

# ❌ WRONG - Deprecated
import gym
env = gym.make("CartPole-v1")

# ✅ CORRECT - Use Gymnasium
import gymnasium as gym
env = gym.make("CartPole-v1")

Gymnasium API (5 Return Values)

import gymnasium as gym

env = gym.make("CartPole-v1")
obs, info = env.reset()

while True:
    action = env.action_space.sample()
    # 5 return values - NOT 4!
    obs, reward, terminated, truncated, info = env.step(action)
    done = terminated or truncated
    if done:
        break

env.close()

TorchRL Overview

TorchRL is PyTorch's official RL library providing:

• Environment wrappers (GymEnv, PettingZooEnv)
• Data collectors (SyncDataCollector, MultiSyncDataCollector)
• Replay buffers with tensor storage
• Loss modules (DQNLoss, ClipPPOLoss)
• Advantage estimation (GAE)

Installation

pip install torchrl gymnasium

TorchRL Environment Wrappers

GymEnv - Single Environment

from torchrl.envs import GymEnv
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create environment
env = GymEnv("CartPole-v1", device=device)

# Get specs
print(f"Observation spec: {env.observation_spec}")
print(f"Action spec: {env.action_spec}")

# Reset and step
td = env.reset()
print(f"Initial state: {td}")

action = env.action_spec.rand()
td = env.step(td.set("action", action))
print(f"Next state: {td}")

Parallel Environments

from torchrl.envs import ParallelEnv, GymEnv

# Run 8 environments in parallel
def make_env():
    return GymEnv("CartPole-v1")

env = ParallelEnv(8, make_env)
td = env.reset()  # Shape: [8, ...]

TransformedEnv - Preprocessing

from torchrl.envs import TransformedEnv, GymEnv
from torchrl.envs.transforms import (
    RewardSum, StepCounter, TransformedEnv,
    ObservationNorm, RewardClipping
)

base_env = GymEnv("CartPole-v1")
env = TransformedEnv(
    base_env,
    transform=Compose(
        ObservationNorm(in_keys=["observation"]),
        RewardClipping(-1, 1),
        StepCounter(),
        RewardSum(),
    )
)

Deep Q-Networks (DQN)

DQN uses a neural network to approximate Q-values with experience replay and target networks for stability.

Q-Network Architecture

import torch
import torch.nn as nn
from tensordict.nn import TensorDictModule

class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )

    def forward(self, x):
        return self.net(x)

# Wrap for TorchRL
qnet = QNetwork(state_dim=4, action_dim=2)
qvalue_module = TensorDictModule(
    qnet,
    in_keys=["observation"],
    out_keys=["action_value"]
)

DQN with TorchRL

import torch
from torchrl.envs import GymEnv
from torchrl.collectors import SyncDataCollector
from torchrl.data import ReplayBuffer, LazyTensorStorage
from torchrl.modules import QValueActor, EGreedyModule
from torchrl.objectives import DQNLoss, HardUpdate
from tensordict.nn import TensorDictModule, TensorDictSequential

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Environment
env = GymEnv("CartPole-v1", device=device)
state_dim = env.observation_spec["observation"].shape[-1]
action_dim = env.action_spec.space.n

# Q-Network
qnet = nn.Sequential(
    nn.Linear(state_dim, 128), nn.ReLU(),
    nn.Linear(128, 128), nn.ReLU(),
    nn.Linear(128, action_dim)
)
qvalue_module = TensorDictModule(
    qnet, in_keys=["observation"], out_keys=["action_value"]
)

# Actor with epsilon-greedy exploration
actor = QValueActor(qvalue_module, spec=env.action_spec)
exploration_module = EGreedyModule(
    spec=env.action_spec,
    eps_init=1.0,
    eps_end=0.05,
    annealing_num_steps=10000
)
policy = TensorDictSequential(actor, exploration_module)

# Replay buffer
replay_buffer = ReplayBuffer(
    storage=LazyTensorStorage(max_size=100_000, device=device),
    batch_size=64
)

# Data collector
collector = SyncDataCollector(
    env,
    policy,
    frames_per_batch=128,
    total_frames=100_000,
    device=device
)

# Loss module
loss_module = DQNLoss(
    value_network=actor,
    loss_function="smooth_l1",
    delay_value=True
)
loss_module.make_value_estimator(gamma=0.99)

# Target network updater
target_updater = HardUpdate(
    loss_module,
    value_network_update_interval=1000
)

# Optimizer
optimizer = torch.optim.Adam(loss_module.parameters(), lr=1e-3)

DQN Training Loop

for i, data in enumerate(collector):
    # Add to replay buffer
    replay_buffer.extend(data)

    if len(replay_buffer) < 1000:
        continue

    # Sample batch
    batch = replay_buffer.sample()

    # Compute loss
    loss_dict = loss_module(batch)
    loss = loss_dict["loss"]

    # Optimize
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(loss_module.parameters(), 1.0)
    optimizer.step()

    # Update target network
    target_updater.step()

    # Update exploration epsilon
    exploration_module.step(data.numel())

    if i % 100 == 0:
        print(f"Step {i}, Loss: {loss.item():.4f}")

Proximal Policy Optimization (PPO)

PPO is an on-policy algorithm using clipped surrogate objective for stable training.

Actor-Critic Architecture

import torch
import torch.nn as nn
from tensordict.nn import TensorDictModule
from torchrl.modules import ProbabilisticActor, ValueOperator
from torch.distributions import Categorical

class ActorCriticNet(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh()
        )
        self.policy_head = nn.Linear(hidden_dim, action_dim)
        self.value_head = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        features = self.shared(x)
        return self.policy_head(features), self.value_head(features)

PPO with TorchRL

import torch
from torchrl.envs import GymEnv, TransformedEnv
from torchrl.envs.transforms import RewardSum, StepCounter
from torchrl.collectors import SyncDataCollector
from torchrl.modules import ProbabilisticActor, ValueOperator
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE
from tensordict.nn import TensorDictModule

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Environment with transforms
base_env = GymEnv("CartPole-v1", device=device)
env = TransformedEnv(base_env, StepCounter())

state_dim = env.observation_spec["observation"].shape[-1]
action_dim = env.action_spec.space.n

# Actor network
actor_net = nn.Sequential(
    nn.Linear(state_dim, 64), nn.Tanh(),
    nn.Linear(64, 64), nn.Tanh(),
    nn.Linear(64, action_dim)
)
actor_module = TensorDictModule(
    actor_net, in_keys=["observation"], out_keys=["logits"]
)
actor = ProbabilisticActor(
    actor_module,
    in_keys=["logits"],
    out_keys=["action"],
    distribution_class=torch.distributions.Categorical,
    return_log_prob=True
)

# Critic network
critic_net = nn.Sequential(
    nn.Linear(state_dim, 64), nn.Tanh(),
    nn.Linear(64, 64), nn.Tanh(),
    nn.Linear(64, 1)
)
critic = ValueOperator(
    TensorDictModule(critic_net, in_keys=["observation"], out_keys=["state_value"])
)

# Move to device
actor = actor.to(device)
critic = critic.to(device)

# Advantage estimation (GAE)
advantage_module = GAE(
    gamma=0.99,
    lmbda=0.95,
    value_network=critic,
    average_gae=True
)

# PPO Loss
loss_module = ClipPPOLoss(
    actor_network=actor,
    critic_network=critic,
    clip_epsilon=0.2,
    entropy_coef=0.01,
    critic_coef=0.5,
    normalize_advantage=True
)

# Data collector
collector = SyncDataCollector(
    env,
    actor,
    frames_per_batch=2048,
    total_frames=100_000,
    device=device
)

# Optimizer
optimizer = torch.optim.Adam(loss_module.parameters(), lr=3e-4)

PPO Training Loop

num_epochs = 10  # PPO epochs per batch

for batch_idx, data in enumerate(collector):
    # Compute advantages
    with torch.no_grad():
        advantage_module(data)

    # PPO update epochs
    for epoch in range(num_epochs):
        # Mini-batch updates
        for minibatch in data.split(256):
            loss_dict = loss_module(minibatch)

            loss = (
                loss_dict["loss_objective"] +
                loss_dict["loss_critic"] +
                loss_dict["loss_entropy"]
            )

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(loss_module.parameters(), 0.5)
            optimizer.step()

    # Logging
    episode_reward = data["next", "reward"].sum().item()
    print(f"Batch {batch_idx}, Reward: {episode_reward:.2f}")

Replay Buffers

Basic Replay Buffer

from torchrl.data import ReplayBuffer, LazyTensorStorage

buffer = ReplayBuffer(
    storage=LazyTensorStorage(max_size=100_000),
    batch_size=64
)

# Add data
buffer.extend(data)

# Sample
batch = buffer.sample()

Prioritized Experience Replay

from torchrl.data import PrioritizedReplayBuffer, LazyTensorStorage

buffer = PrioritizedReplayBuffer(
    storage=LazyTensorStorage(max_size=100_000),
    alpha=0.6,
    beta=0.4,
    batch_size=64
)

Multi-Agent RL with PettingZoo

from torchrl.envs import PettingZooEnv

# Cooperative/competitive environments
env = PettingZooEnv(
    env="simple_spread_v3",
    parallel=True
)

td = env.reset()
print(f"Agents: {env.agents}")

Distributed Training

Parallel Data Collection

from torchrl.collectors import MultiSyncDataCollector

collector = MultiSyncDataCollector(
    [make_env] * 4,  # 4 parallel collectors
    policy,
    frames_per_batch=512,
    total_frames=100_000,
    device=device
)

Multi-GPU with torchrun

# Launch distributed PPO
torchrun --nproc_per_node=4 train_ppo.py

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
local_rank = int(os.environ["LOCAL_RANK"])
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

Key Hyperparameters

DQN Hyperparameters

PPO Hyperparameters

Best Practices

1. Always use Gymnasium - OpenAI Gym is deprecated 2. Start simple - CartPole before Atari 3. Monitor training - Use TensorBoard or W&B 4. Tune one thing at a time - Systematic hyperparameter search 5. Use vectorized environments - Faster data collection 6. Clip gradients - Prevents training instability 7. Normalize observations - Helps neural network learning 8. Use torch.compile - Faster training (CUDA only)

Common Issues

NaN Loss

• Check for division by zero in advantage normalization
• Verify reward scaling is reasonable
• Ensure learning rate isn't too high

No Learning

• Increase exploration (epsilon/entropy)
• Check environment reward structure
• Verify network is receiving gradients

Unstable Training

• Reduce learning rate
• Increase batch size
• Adjust clip range (PPO)
• Use target network (DQN)