> ## Documentation Index
> Fetch the complete documentation index at: https://docs.monostate.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Módulo de Treinamento RL

> Aprendizado por reforço avançado para LLMs

# Módulo de Treinamento RL

O AITraining inclui um módulo abrangente de aprendizado por reforço para cenários avançados de treinamento de LLM.

<Note>
  Os comandos CLI (`--trainer ppo`, `--trainer dpo`, `--trainer reward`) usam implementações da biblioteca TRL para estabilidade. O módulo `autotrain.trainers.rl` documentado aqui fornece blocos de construção de nível inferior para pipelines de treinamento RL personalizados.
</Note>

## Visão Geral

O módulo RL fornece:

* **PPO Trainer** - Proximal Policy Optimization com penalidade KL e GAE
* **DPO Trainer** - Otimização Direta de Preferências a partir de dados de preferência
* **Modelos de Recompensa** - Modelos de recompensa padrão, pareados e multi-objetivo
* **Ambientes RL** - Ambientes de geração de texto, matemática, código e comparação de preferências
* **Pipeline Assíncrono** - Treinamento forward-backward com acúmulo de gradiente

## Treinamento PPO

### Configuração

```python theme={null}
from autotrain.trainers.rl import PPOConfig, PPOTrainer

config = PPOConfig(
    model_name="google/gemma-2-2b",
    learning_rate=1e-5,
    batch_size=16,
    mini_batch_size=4,
    gradient_accumulation_steps=1,

    # PPO hyperparameters
    ppo_epochs=4,
    gamma=0.99,           # Discount factor
    lam=0.95,             # GAE lambda
    clip_ratio=0.2,       # PPO clip ratio
    value_clip=0.2,       # Value function clip
    max_grad_norm=1.0,    # Gradient clipping

    # KL penalty
    kl_penalty_coef=0.01,
    kl_target=0.01,
    kl_horizon=10000,     # Horizon for adaptive KL

    # Coefficients
    entropy_coef=0.01,    # Entropy regularization
    value_coef=0.5,       # Value function coefficient

    # Generation
    max_new_tokens=128,
    temperature=0.7,
    top_p=0.9,

    # Training loop
    num_iterations=100,
    save_every=10,
    eval_every=5,
    device=None,          # Auto-detected
)
```

### Arquitetura PPO

A implementação PPO usa um wrapper `PPOModel` que adiciona um cabeçalho de valor a qualquer LM causal:

```python theme={null}
# PPOModel wraps base model with ValueHead
class PPOModel(nn.Module):
    def __init__(self, base_model):
        self.base_model = base_model
        self.value_head = ValueHead(hidden_size)

# ValueHead architecture
class ValueHead(nn.Module):
    # hidden -> ReLU -> output (scalar value)
```

### Controlador KL Adaptativo

O `AdaptiveKLController` ajusta automaticamente o coeficiente de penalidade KL para manter a divergência KL próxima ao alvo:

```python theme={null}
# Automatically managed by PPOTrainer
# Adjusts kl_penalty_coef based on current KL vs target
```

### Loop de Treinamento

```python theme={null}
# Initialize trainer with custom reward function
def my_reward_fn(prompts, responses, metadata=None):
    rewards = []
    for response in responses:
        score = evaluate_response(response)
        rewards.append(score)
    return rewards

trainer = PPOTrainer(
    config=config,
    tokenizer=tokenizer,  # Optional, loaded from model if not provided
    reward_fn=my_reward_fn,
)

# Train on prompts
prompts = ["Write a poem about...", "Explain quantum..."]
metrics = trainer.train(prompts, num_iterations=100)
```

### Recursos Principais

| Recurso                  | Descrição                                                                          |
| ------------------------ | ---------------------------------------------------------------------------------- |
| Adaptive KL Controller   | Ajusta automaticamente o coeficiente de penalidade KL com base na KL atual vs alvo |
| GAE Advantage Estimation | Estimativa de Vantagem Generalizada para treinamento estável                       |
| Value Head               | Função de valor separada para crítico (PPOModel envolve modelo base)               |
| Reference Model          | Cópia congelada para prevenir deriva                                               |
| Async Training           | Usa AsyncTrainingClient para forward-backward eficiente                            |

## Treinamento DPO

Treine diretamente a partir de dados de preferência sem um modelo de recompensa separado.

### Configuração

```python theme={null}
from autotrain.trainers.rl import DPOConfig, DPOTrainer

config = DPOConfig(
    model_name="google/gemma-2-2b",
    learning_rate=1e-6,
    batch_size=8,
    gradient_accumulation_steps=2,

    # DPO hyperparameters
    beta=0.1,              # Temperature parameter
    label_smoothing=0.0,   # For robustness
    reference_free=False,  # Use reference model

    # Training
    num_epochs=1,
    max_grad_norm=1.0,
    warmup_ratio=0.1,

    # Sequence lengths
    max_length=512,
    max_prompt_length=256,

    # Checkpointing
    eval_every=100,
    save_every=500,
    device=None,           # Auto-detected
)
```

### Dataset de Preferência

```python theme={null}
from autotrain.trainers.rl.dpo import PreferenceDataset

# Create dataset from preference pairs
dataset = PreferenceDataset(
    prompts=["What is AI?", "Explain gravity"],
    chosen=["AI is artificial intelligence...", "Gravity is a force..."],
    rejected=["idk lol", "its like magnets"],
    tokenizer=tokenizer,
    max_length=512,
    max_prompt_length=256,
)

# Train
trainer = DPOTrainer(config=config, tokenizer=tokenizer)
metrics = trainer.train(dataset, eval_dataset=eval_dataset)
```

<Note>
  `PreferenceDataset` deve ser importado diretamente de `autotrain.trainers.rl.dpo` pois não é exportado no `__init__.py` principal.
</Note>

### DPO Sem Referência

Para treinamento sem modelo de referência:

```python theme={null}
config = DPOConfig(
    model_name="google/gemma-2-2b",
    reference_free=True,  # No reference model needed
    beta=0.1,
)
```

## Modelos de Recompensa

### Modelo de Recompensa Padrão

```python theme={null}
from autotrain.trainers.rl import RewardModel, RewardModelConfig, RewardModelTrainer

config = RewardModelConfig(
    model_name="bert-base-uncased",
    num_labels=1,
    pooling_strategy="last",  # "mean", "last", or "cls"
    dropout_prob=0.1,
    temperature=1.0,          # Temperature scaling for rewards

    # LoRA settings
    use_lora=True,
    lora_rank=8,
    lora_alpha=16,
    lora_dropout=0.1,

    # Training
    learning_rate=1e-4,
    warmup_steps=100,
    gradient_accumulation_steps=1,
)

model = RewardModel(config)
```

### Treinamento em Preferências

```python theme={null}
trainer = RewardModelTrainer(
    model=model,
    tokenizer=tokenizer,
    config=config,
    device=None,  # Auto-detected
)

trainer.train_on_preferences(
    chosen_texts=["Good response 1", "Good response 2"],
    rejected_texts=["Bad response 1", "Bad response 2"],
    num_epochs=3,
    batch_size=8,
)

# Save/load
trainer.save_model("reward_model.pt")
trainer.load_model("reward_model.pt")
```

### Modelo de Recompensa Pareado

Para comparação direta de preferências usando modelo Bradley-Terry:

```python theme={null}
from autotrain.trainers.rl import PairwiseRewardModel

model = PairwiseRewardModel(config)

# Forward pass compares two inputs
preference_score = model.forward_pair(
    input_ids_a, attention_mask_a,
    input_ids_b, attention_mask_b,
)

# Bradley-Terry loss for training
loss = model.compute_bradley_terry_loss(
    input_ids_a, attention_mask_a,
    input_ids_b, attention_mask_b,
    labels,  # 1 if A preferred, 0 if B preferred
)
```

### Modelo de Recompensa Multi-Objetivo

Combine múltiplos sinais de recompensa:

```python theme={null}
from autotrain.trainers.rl import MultiObjectiveRewardModel

model = MultiObjectiveRewardModel(
    config=config,
    num_objectives=3,
    objective_weights=[0.5, 0.3, 0.2],  # Helpfulness, safety, honesty
)

# Get all objectives
outputs = model(input_ids, attention_mask, return_all_objectives=True)
# outputs["rewards"] shape: (batch_size, 3)
# outputs["combined_reward"] shape: (batch_size, 1)

# Multi-objective loss
loss, per_objective_losses = model.compute_multi_objective_loss(
    input_ids, attention_mask,
    target_rewards,      # Shape: (batch_size, num_objectives)
    objective_mask=None, # Optional: which objectives to train
)
```

## Ambientes RL

### Dataclasses de Ambiente

```python theme={null}
from autotrain.trainers.rl.environments import Observation, StepResult, Trajectory

# Observation from environment
@dataclass
class Observation:
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    prompt: str
    metadata: Dict[str, Any]

# Result from env.step()
@dataclass
class StepResult:
    reward: float
    done: bool
    next_observation: Optional[Observation]
    info: Dict[str, Any]
    metrics: Dict[str, float]

# Full episode trajectory
@dataclass
class Trajectory:
    observations: List[Observation]
    actions: List[torch.Tensor]
    rewards: List[float]
    logprobs: List[torch.Tensor]
    done: bool
    total_reward: float
    metrics: Dict[str, Any]
```

### Ambiente de Geração de Texto

```python theme={null}
from autotrain.trainers.rl import TextGenerationEnv

env = TextGenerationEnv(
    tokenizer=tokenizer,
    prompts=["Write a story about...", "Explain how..."],
    max_length=512,
    reward_fn=my_reward_function,  # Optional, default is length-based
    stop_sequences=["</s>", "\n\n"],
    temperature=1.0,
)

# Reset and step
observation = env.reset()
result = env.step(action_token)
# result.reward, result.done, result.next_observation

# Render current state
print(env.render())
```

### Ambiente Multi-Objetivo

```python theme={null}
from autotrain.trainers.rl import MultiObjectiveRewardEnv

def correctness_reward(prompt, generated, full_text):
    return 1.0 if is_correct(generated) else 0.0

def formatting_reward(prompt, generated, full_text):
    return 0.5 if properly_formatted(generated) else 0.0

env = MultiObjectiveRewardEnv(
    tokenizer=tokenizer,
    prompts=prompts,
    reward_components={
        "correctness": correctness_reward,
        "formatting": formatting_reward,
    },
    reward_weights={
        "correctness": 1.0,
        "formatting": 0.1,
    },
)

# Step returns component rewards in metrics
result = env.step(action)
# result.metrics["reward_correctness"], result.metrics["reward_formatting"]
```

### Ambiente de Comparação de Preferências

Para coleta de dados RLHF e DPO:

```python theme={null}
from autotrain.trainers.rl import PreferenceComparisonEnv

env = PreferenceComparisonEnv(
    tokenizer=tokenizer,
    prompts=prompts,
    preference_model=preference_model,  # Optional
    human_feedback_fn=feedback_fn,      # Optional callback
    max_length=512,
)

# Generates pairs of responses and computes preference
observation = env.reset()
result1 = env.step(response1_tokens)  # First response
result2 = env.step(response2_tokens)  # Second response, computes preference
```

### Ambientes Integrados

```python theme={null}
from autotrain.trainers.rl import create_math_problem_env, create_code_generation_env

# Math problem solving (correctness + formatting rewards)
math_env = create_math_problem_env(tokenizer)

# Code generation (syntax + style rewards)
code_env = create_code_generation_env(tokenizer)
```

## Pipeline Forward-Backward

Treinamento assíncrono com acúmulo de gradiente:

```python theme={null}
from autotrain.trainers.rl import ForwardBackwardPipeline

# Low-level pipeline
pipeline = ForwardBackwardPipeline(
    model=model,
    device="cuda",
    max_workers=2,                    # Thread pool size
    gradient_accumulation_steps=4,
)

# Queue forward-backward pass
future = pipeline.forward_backward(
    input_ids=input_ids,
    attention_mask=attention_mask,
    labels=labels,
    loss_fn="cross_entropy",
)

# Get result (blocks until complete)
result = future.result()
print(f"Loss: {result.loss}")

# Queue optimizer step
optim_future = pipeline.optim_step(
    optimizer=optimizer,
    scheduler=scheduler,  # Optional
    max_grad_norm=1.0,
)
optim_result = optim_future.result()
```

### Funções de Perda Integradas

O pipeline suporta várias funções de perda integradas:

| Função de Perda         | Descrição                              | kwargs Obrigatórios                                              |
| ----------------------- | -------------------------------------- | ---------------------------------------------------------------- |
| `"cross_entropy"`       | Perda padrão de modelagem de linguagem | None                                                             |
| `"importance_sampling"` | RL com amostragem de importância       | `old_logprobs`, `advantages`                                     |
| `"ppo"`                 | Perda PPO completa                     | `old_log_probs`, `advantages`, opcionalmente `values`, `returns` |

### Funções de Perda Personalizadas

```python theme={null}
def custom_loss_fn(model, inputs, outputs, **kwargs):
    # Your custom loss computation
    logits = outputs.logits
    # ... compute loss ...
    return loss_tensor  # Must be scalar

future = pipeline.forward_backward_custom(
    input_ids=input_ids,
    custom_loss_fn=custom_loss_fn,
    attention_mask=attention_mask,  # Optional
    my_param=42,  # Passed to loss function via kwargs
)
```

### Cliente de Alto Nível

```python theme={null}
from autotrain.trainers.rl.forward_backward import AsyncTrainingClient

client = AsyncTrainingClient(
    model=model,
    reference_model=reference_model,  # For PPO/DPO
    device="cuda",
    gradient_accumulation_steps=4,
)

# Training step
fwd_future = client.forward_backward(batch, loss_fn="cross_entropy")
optim_future = client.optim_step(optimizer, max_grad_norm=1.0)

# Forward only (for reference model)
ref_future = client.forward(batch, use_reference=True)

# Clean up
client.shutdown()
```

<Note>
  `AsyncTrainingClient` deve ser importado diretamente de `autotrain.trainers.rl.forward_backward` pois não é exportado no `__init__.py` principal.
</Note>

## Checkpointing

```python theme={null}
# Save checkpoint
checkpoint_info = pipeline.save_state("checkpoint_1000")
# Returns: {"path": ..., "model_path": ..., "optimizer_path": ..., "state_path": ...}

# Load checkpoint
pipeline.load_state("checkpoints/checkpoint_1000")
```

## Amostragem

Gere amostras durante o treinamento:

```python theme={null}
samples = pipeline.sample(
    prompt=prompt_tokens,  # List[int] or Tensor
    max_tokens=100,
    temperature=0.7,
    top_k=50,
    top_p=0.9,
    stop=[tokenizer.eos_token_id],
)

print(f"Generated: {samples['tokens']}")
print(f"Logprobs: {samples['logprobs']}")
print(f"Prompt: {samples['prompt']}")
```

## Melhores Práticas

### Treinamento PPO

1. **Comece com coeficiente KL pequeno** - Deixe o controlador adaptativo ajustar
2. **Use acúmulo de gradiente** - Tamanhos de batch efetivos maiores são mais estáveis
3. **Monitore divergência KL** - Deve permanecer próximo ao alvo
4. **Aqueça a função de valor** - Treine o crítico antes do PPO completo

### Treinamento DPO

1. **Dados de preferência de alta qualidade** - Qualidade importa mais que quantidade
2. **Taxa de aprendizado baixa** - 1e-6 a 1e-5 recomendado
3. **Suavização de rótulo** - 0.1 pode melhorar robustez
4. **Avalie frequentemente** - Acompanhe precisão e margem de recompensa

### Modelagem de Recompensa

1. **Dados balanceados** - Exemplos escolhido/rejeitado iguais
2. **Prompts diversos** - Cubra casos de uso esperados
3. **LoRA para eficiência** - Fine-tuning eficiente de modelos grandes
4. **Multi-objetivo** - Separe sinais de segurança e utilidade

## Integração CLI

Para uso em produção, o CLI fornece interfaces mais simples usando implementações TRL:

```bash theme={null}
# PPO training (uses TRL PPOTrainer)
aitraining llm --train \
  --model google/gemma-2-2b \
  --trainer ppo \
  --reward-model ./my-reward-model

# DPO training (uses TRL DPOTrainer)
aitraining llm --train \
  --model google/gemma-2-2b \
  --trainer dpo \
  --dpo-beta 0.1

# Reward model training
aitraining llm --train \
  --model google/gemma-2-2b \
  --trainer reward \
  --data-path ./preference_data.jsonl
```

## Próximos Passos

<CardGroup cols={2}>
  <Card title="PPO Training" href="/advanced/ppo-rl-training" icon="brain">
    Guia CLI para PPO
  </Card>

  <Card title="DPO Training" href="/advanced/dpo-training" icon="scale-balanced">
    Guia CLI para DPO
  </Card>

  <Card title="Reward Modeling" href="/advanced/reward-modeling" icon="star">
    Treinar modelos de recompensa
  </Card>

  <Card title="ORPO Training" href="/advanced/orpo-training" icon="chart-line">
    Otimização de Preferências por Razão de Chances
  </Card>
</CardGroup>