### Training Workflow with Hydra

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Commands to install dependencies, prepare datasets, and launch training jobs with various environment configurations.

```bash
uv venv --python=3.10
source .venv/bin/activate
uv pip install stable-worldmodel[train,env]
export STABLEWM_HOME=/path/to/your/storage
tar --zstd -xvf pusht_expert_train.tar.zst
python train.py data=pusht
python train.py data=pusht --trainer.max_epochs=200 --optimizer.lr=1e-4 --loss.sigreg.weight=0.05
```

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

### GET /policy/AutoCostModel

Source: https://github.com/lucas-maes/le-wm/blob/main/README.md

Initializes a cost model for Model Predictive Control (MPC) by loading a checkpoint from the local environment.

```APIDOC
## GET /policy/AutoCostModel

### Description
Initializes the AutoCostModel by loading a pre-trained checkpoint. This is primarily used for MPC tasks.

### Method
GET (Internal API call)

### Endpoint
swm.policy.AutoCostModel(run_name, cache_dir=None)

### Parameters
#### Path Parameters
- **run_name** (string) - Required - The checkpoint path relative to $STABLEWM_HOME, excluding the '_object.ckpt' suffix.

#### Query Parameters
- **cache_dir** (string) - Optional - Override for the default checkpoint root directory ($STABLEWM_HOME).

### Request Example
```python
import stable_worldmodel as swm
cost = swm.policy.AutoCostModel('pusht/lewm')
```

### Response
#### Success Response (200)
- **model** (object) - Returns a module instance in eval mode with accessible PyTorch weights via .state_dict().
```

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

### Load Pretrained Checkpoints for Inference and Planning

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Demonstrates how to load pretrained checkpoints using the stable_worldmodel API for tasks like MPC planning. It covers loading cost models, configuring planning parameters, creating CEM solvers, and accessing underlying PyTorch weights. Dependencies include stable_worldmodel and torch.

```python
import stable_worldmodel as swm
import torch

# Load cost model for MPC planning
# Path is relative to $STABLEWM_HOME, without _object.ckpt suffix
cost_model = swm.policy.AutoCostModel("pusht/lewm")
cost_model = cost_model.to("cuda")
cost_model = cost_model.eval()
cost_model.requires_grad_(False)

# Configure planning
config = swm.PlanConfig(
    horizon=5,
    receding_horizon=5,
    action_block=5,
)

# Create CEM solver
solver = swm.solver.CEMSolver(
    model=cost_model,
    batch_size=1,
    num_samples=300,
    var_scale=1.0,
    n_steps=30,
    topk=30,
    device="cuda",
)

# Create policy
policy = swm.policy.WorldModelPolicy(
    solver=solver,
    config=config,
    process={},
    transform={},
)

# Access underlying PyTorch weights
state_dict = cost_model.state_dict()
print(f"Model has {sum(p.numel() for p in cost_model.parameters())} parameters")
```

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

### Initialize LeWM JEPA Architecture

Source: https://context7.com/lucas-maes/le-wm/llms.txt

This snippet demonstrates how to instantiate the core components of the LeWM architecture, including the Vision Transformer encoder, autoregressive predictor, action encoder, and MLP projectors. It shows the assembly of these modules into a functional JEPA world model and performs a basic forward pass for encoding and prediction.

```python
import torch
from jepa import JEPA
from module import ARPredictor, Embedder, MLP
import stable_pretraining as spt

encoder = spt.backbone.utils.vit_hf("tiny", patch_size=14, image_size=224, pretrained=False, use_mask_token=False)
hidden_dim = encoder.config.hidden_size
embed_dim = 192
action_dim = 2

predictor = ARPredictor(num_frames=3, input_dim=embed_dim, hidden_dim=hidden_dim, output_dim=hidden_dim, depth=6, heads=16, mlp_dim=2048, dim_head=64, dropout=0.1)
action_encoder = Embedder(input_dim=action_dim * 5, emb_dim=embed_dim)
projector = MLP(input_dim=hidden_dim, output_dim=embed_dim, hidden_dim=2048, norm_fn=torch.nn.BatchNorm1d)
pred_proj = MLP(input_dim=hidden_dim, output_dim=embed_dim, hidden_dim=2048, norm_fn=torch.nn.BatchNorm1d)

world_model = JEPA(encoder=encoder, predictor=predictor, action_encoder=action_encoder, projector=projector, pred_proj=pred_proj)

batch = {"pixels": torch.randn(4, 3, 3, 224, 224), "action": torch.randn(4, 3, 10)}
output = world_model.encode(batch)
pred_emb = world_model.predict(output["emb"][:, :3], output["act_emb"][:, :3])
```

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

### Data Preprocessing Utilities for Training and Evaluation

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Provides utility functions for image preprocessing and column normalization, preparing data for training and evaluation. It shows how to set up image transformations, load HDF5 datasets, create normalizers for specific columns (e.g., action), and compose transforms. Dependencies include torch, numpy, stable_worldmodel, and stable_pretraining.data.transforms.

```python
import torch
import numpy as np
from utils import get_img_preprocessor, get_column_normalizer
import stable_worldmodel as swm

# Image preprocessing pipeline
img_transform = get_img_preprocessor(
    source="pixels",
    target="pixels",
    img_size=224,
)

# Load dataset
dataset = swm.data.HDF5Dataset(
    "pusht_expert_train",
    keys_to_load=["pixels", "action", "proprio", "state"],
    keys_to_cache=["action", "proprio", "state"],
)

# Create normalizer for action column
action_normalizer = get_column_normalizer(
    dataset,
    source="action",
    target="action",
)

# Compose transforms
from stable_pretraining.data.transforms import Compose
transform = Compose(img_transform, action_normalizer)
dataset.transform = transform

# Access a sample
sample = dataset[0]
# sample["pixels"]: normalized image tensor
# sample["action"]: normalized action tensor
```

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

### Initialize and Run ARPredictor

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Sets up the transformer-based autoregressive predictor to forecast next-state embeddings conditioned on action inputs.

```python
import torch
from module import ARPredictor
predictor = ARPredictor(num_frames=3, input_dim=192, hidden_dim=768, output_dim=768, depth=6, heads=16, mlp_dim=2048, dim_head=64, dropout=0.1, emb_dropout=0.0)
state_emb = torch.randn(4, 3, 192)
action_cond = torch.randn(4, 3, 192)
pred_emb = predictor(state_emb, action_cond)
print(f"Input shape: {state_emb.shape}")
print(f"Output shape: {pred_emb.shape}")
```

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

### Perform Model Predictive Control with CEM

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Demonstrates generating action candidates and selecting the best sequence based on cost minimization using a world model.

```python
num_samples = 300
horizon = 10
action_dim = 10
action_candidates = torch.randn(1, num_samples, horizon, action_dim).to(device)
with torch.no_grad():
    costs = world_model.get_cost(info, action_candidates)
best_idx = costs.argmin(dim=1)
best_action = action_candidates[0, best_idx]
print(f"Best action sequence shape: {best_action.shape}")
print(f"Minimum cost: {costs.min().item():.4f}")
```

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

### Load Object Checkpoint with stable_worldmodel API

Source: https://github.com/lucas-maes/le-wm/blob/main/README.md

Demonstrates loading a cost model checkpoint using the stable_worldmodel API. This method is suitable for direct use with the API or for evaluation scripts. It requires the 'stable_worldmodel' library and accepts a run name relative to the STABLEWM_HOME environment variable.

```python
import stable_worldmodel as swm

# Load the cost model (for MPC)
cost = swm.policy.AutoCostModel('pusht/lewm')
```

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

### Autoregressive Rollout for Planning

Source: https://context7.com/lucas-maes/le-wm/llms.txt

This section explains how to use the `rollout` method for autoregressive prediction over action sequences, enabling model predictive control.

```APIDOC
## Autoregressive Rollout for Planning

### Description
The `rollout` method performs autoregressive prediction over action sequences for model predictive control. It takes initial observations and candidate action sequences, then predicts future state embeddings by iteratively applying the predictor.

### Method
```python
import torch

# Assume world_model is a trained JEPA instance
world_model.eval()

# Initial observation info dict
info = {
    "pixels": torch.randn(1, 1, 3, 3, 224, 224),  # (B, S=1, T=3, C, H, W)
}

# Candidate action sequences: (B, S, T, action_dim)
# B=1 batch, S=100 samples, T=10 timesteps, 10-dim actions
action_candidates = torch.randn(1, 100, 10, 10)

# Perform rollout
with torch.no_grad():
    info = world_model.rollout(
        info,
        action_candidates,
        history_size=3,  # number of past frames to consider
    )

# info["predicted_emb"]: (1, 100, 11, 192)
# Contains predicted embeddings for all action candidates over time

# The rollout enables MPC by evaluating which action sequence
# leads to states closest to the goal
print(f"Predicted embeddings shape: {info['predicted_emb'].shape}")
```
```

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

### Compute MPC Costs

Source: https://context7.com/lucas-maes/le-wm/llms.txt

This snippet shows how to prepare the input dictionary for cost computation in a model predictive control context. It initializes the necessary tensors on the GPU to evaluate the distance between predicted final states and goal states.

```python
import torch

world_model.eval()
device = "cuda"

info = {
    "pixels": torch.randn(1, 1, 3, 3, 224, 224).to(device),
    "goal": torch.randn(1, 1, 3, 224, 224).to(device),
    "action": torch.randn(1, 1, 3, 10).to(device)
}
```

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

### JEPA Core Architecture

Source: https://context7.com/lucas-maes/le-wm/llms.txt

This section details the construction and basic usage of the JEPA core architecture, including encoding observations and predicting future states.

```APIDOC
## JEPA Core Architecture

### Description
The JEPA class implements the main world model architecture, combining an encoder for pixel observations, an action encoder, and a predictor for next-state embedding prediction. It supports encoding observations, predicting future states, and computing costs for model predictive control.

### Method
```python
import torch
from jepa import JEPA
from module import ARPredictor, Embedder, MLP
import stable_pretraining as spt

# Create encoder (Vision Transformer)
encoder = spt.backbone.utils.vit_hf(
    "tiny",
    patch_size=14,
    image_size=224,
    pretrained=False,
    use_mask_token=False,
)

hidden_dim = encoder.config.hidden_size
embed_dim = 192
action_dim = 2  # e.g., for PushT environment

# Create predictor (autoregressive transformer)
predictor = ARPredictor(
    num_frames=3,  # history size
    input_dim=embed_dim,
    hidden_dim=hidden_dim,
    output_dim=hidden_dim,
    depth=6,
    heads=16,
    mlp_dim=2048,
    dim_head=64,
    dropout=0.1,
)

# Create action encoder
action_encoder = Embedder(input_dim=action_dim * 5, emb_dim=embed_dim)  # action_dim * frameskip

# Create projectors
projector = MLP(
    input_dim=hidden_dim,
    output_dim=embed_dim,
    hidden_dim=2048,
    norm_fn=torch.nn.BatchNorm1d,
)
pred_proj = MLP(
    input_dim=hidden_dim,
    output_dim=embed_dim,
    hidden_dim=2048,
    norm_fn=torch.nn.BatchNorm1d,
)

# Assemble JEPA world model
world_model = JEPA(
    encoder=encoder,
    predictor=predictor,
    action_encoder=action_encoder,
    projector=projector,
    pred_proj=pred_proj,
)

# Example: Encode observations and actions
batch = {
    "pixels": torch.randn(4, 3, 3, 224, 224),  # (B, T, C, H, W)
    "action": torch.randn(4, 3, 10),            # (B, T, action_dim*frameskip)
}
output = world_model.encode(batch)
# output["emb"]: (4, 3, 192) - pixel embeddings
# output["act_emb"]: (4, 3, 192) - action embeddings

# Predict next state
pred_emb = world_model.predict(output["emb"][:, :3], output["act_emb"][:, :3])
# pred_emb: (4, 3, 192) - predicted next embeddings
```
```

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

### Cost Computation for Model Predictive Control

Source: https://context7.com/lucas-maes/le-wm/llms.txt

This section covers the `get_cost` method, used to evaluate action candidates by computing costs between predicted and goal states for optimal action selection.

```APIDOC
## Cost Computation for Model Predictive Control

### Description
The `get_cost` method evaluates action candidates by computing the cost between predicted final states and goal states. This enables selection of optimal actions for reaching a specified goal.

### Method
```python
import torch

# Assume world_model is a trained JEPA instance on GPU
world_model.eval()
device = "cuda"

# Prepare info dict with current observation and goal
info = {
    "pixels": torch.randn(1, 1, 3, 3, 224, 224).to(device),    # current state
    "goal": torch.randn(1, 1, 3, 224, 224).to(device),         # goal image
    "action": torch.randn(1, 1, 3, 10).to(device),             # initial actions
}

# The actual call to get_cost would follow, likely involving
# a rollout first to get predicted embeddings.
# Example placeholder:
# predicted_embeddings = world_model.rollout(info, action_candidates, history_size=3)
# costs = world_model.get_cost(predicted_embeddings, info["goal"])
```
```

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

### Initialize and Use SIGReg Regularizer

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Configures the Sketch Isotropic Gaussian Regularizer to enforce Gaussian-distributed latent embeddings, preventing representation collapse.

```python
import torch
from module import SIGReg
sigreg = SIGReg(knots=17, num_proj=1024)
embeddings = torch.randn(10, 32, 192)
sigreg_loss = sigreg(embeddings)
print(f"SIGReg loss: {sigreg_loss.item():.4f}")
pred_loss = torch.tensor(0.5)
lambd = 0.09
total_loss = pred_loss + lambd * sigreg_loss
```

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

### Perform Autoregressive Rollout for Planning

Source: https://context7.com/lucas-maes/le-wm/llms.txt

This snippet illustrates the use of the rollout method to perform autoregressive prediction over action sequences. It is designed for model predictive control (MPC) by iteratively applying the predictor to candidate action sequences to forecast future state embeddings.

```python
import torch

world_model.eval()
info = {"pixels": torch.randn(1, 1, 3, 3, 224, 224)}
action_candidates = torch.randn(1, 100, 10, 10)

with torch.no_grad():
    info = world_model.rollout(info, action_candidates, history_size=3)

print(f"Predicted embeddings shape: {info['predicted_emb'].shape}")
```

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

### Evaluate Trained Models

Source: https://context7.com/lucas-maes/le-wm/llms.txt

Executes evaluation scripts using different solvers like CEM or Adam to assess model performance on goal-reaching tasks.

```bash
python eval.py --config-name=pusht.yaml policy=pusht/lewm
python eval.py --config-name=pusht.yaml policy=pusht/lewm solver=adam
python eval.py --config-name=pusht.yaml policy=pusht/lewm eval.num_eval=100 eval.eval_budget=100 plan_config.horizon=10
```

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