### Complete Training Example for ScalableVIN

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

Demonstrates an end-to-end training loop for maze navigation using ScalableVIN. Includes model configuration, data preparation, and optimization. Use for practical application of the library.

```python
import torch
import torch.optim as optim
from scaling_vin_pytorch import ScalableVIN

# Model configuration for maze navigation
model = ScalableVIN(
    state_dim=3,                  # Maze state (walls, goal, agent)
    reward_dim=1,                 # Goal reward signal
    num_actions=4,                # Up, down, left, right
    num_plans=2,                  # Multiple planning strategies
    depth=200,                    # Deep planning for complex mazes
    loss_every_num_layers=8,      # Multi-layer loss
    checkpoint=True,              # Memory efficient
    vi_module_type='dynamic',     # Adaptive transitions
    soft_maxpool=True,
    soft_maxpool_temperature=0.5
)

optimizer = optim.Adam(model.parameters(), lr=1e-4)

# Training loop
model.train()
for epoch in range(100):
    # Simulated maze data
    batch_size = 8
    maze_size = 64

    state = torch.randn(batch_size, 3, maze_size, maze_size)
    reward = torch.randn(batch_size, 1, maze_size, maze_size)
    positions = torch.randint(0, maze_size, (batch_size, 2))
    actions = torch.randint(0, 4, (batch_size,))

    optimizer.zero_grad()
    loss = model(state, reward, positions, actions)
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# Inference
model.eval()
with torch.no_grad():
    test_state = torch.randn(1, 3, maze_size, maze_size)
    test_reward = torch.randn(1, 1, maze_size, maze_size)
    test_pos = torch.randint(0, maze_size, (1, 2))

    logits = model(test_state, test_reward, test_pos)
    action = logits[:, 0].argmax(dim=-1)
    print(f"Predicted action: {action.item()}")
```

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

### Initialize and use ScalableVIN

Source: https://github.com/lucidrains/scaling-vin-pytorch/blob/main/README.md

Instantiate the ScalableVIN model and perform forward passes for training and inference.

```python
import torch
from scaling_vin_pytorch import ScalableVIN

scalable_vin = ScalableVIN(
    state_dim = 3,
    reward_dim = 2,
    num_actions = 10
)

state = torch.randn(2, 3, 32, 32)
reward = torch.randn(2, 2, 32, 32)

agent_positions = torch.randint(0, 32, (2, 2))

target_actions = torch.randint(0, 10, (2,))

loss = scalable_vin(
    state,
    reward,
    agent_positions,
    target_actions
)

loss.backward()

action_logits = scalable_vin(
    state,
    reward,
    agent_positions
)
```

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

### ValueIteration - Invariant Transition Module

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

Standard value iteration using fixed convolutional kernels for state transitions. Best for environments with consistent transition dynamics.

```APIDOC
## ValueIteration - Invariant Transition Module

### Description
Standard value iteration using fixed convolutional kernels for state transitions. Best for environments with consistent transition dynamics.

### Method
POST (or similar, as it's a module initialization and forward pass)

### Endpoint
`/api/value_iteration/invariant` (Hypothetical endpoint for module usage)

### Parameters
#### Initialization Parameters
- **action_channels** (int) - Number of action channels
- **num_plans** (int) - Number of parallel value plans
- **receptive_field** (int) - Kernel size for transitions (must be odd)
- **pad_value** (float) - Padding value for boundaries
- **soft_maxpool** (bool) - Use hard max for action selection
- **soft_maxpool_temperature** (float) - Temperature for soft max-pooling

#### Input Data (for forward pass)
- **values** (torch.Tensor) - Current value estimates (batch_size, num_plans, height, width)
- **rewards** (torch.Tensor) - Reward signal (batch_size, num_plans, height, width)

### Request Example (Initialization)
```python
import torch
from scaling_vin_pytorch import ValueIteration

vi_module = ValueIteration(
    action_channels=8,
    num_plans=1,
    receptive_field=3,
    pad_value=0.0,
    soft_maxpool=False,
    soft_maxpool_temperature=1.0
)
```

### Request Example (Forward Pass)
```python
# Input tensors
batch_size = 4
height, width = 16, 16
num_plans = 1

values = torch.randn(batch_size, num_plans, height, width)
rewards = torch.randn(batch_size, num_plans, height, width)

# Perform one step of value iteration
next_values = vi_module(values, rewards)
```

### Response
#### Success Response (200)
- **next_values** (torch.Tensor) - Updated value estimates after one iteration. Shape: (batch_size, num_plans, height, width)

#### Response Example
```json
{
  "next_values": [[[...]]]
}
```
```

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

### AttentionValueIteration Module

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

Implements attention-based value iteration for computing transitions. Use when a more expressive transition model is needed. Requires PyTorch.

```python
import torch
from scaling_vin_pytorch import AttentionValueIteration

# Create attention-based value iteration module
attn_vi = AttentionValueIteration(
    action_channels=8,
    dim_qk=8,                   # Query/key dimension for attention
    num_plans=1,
    receptive_field=3,          # Attention receptive field
    pad_value=0.0,
    soft_maxpool=False,
    soft_maxpool_temperature=1.0
)

batch_size = 4
height, width = 16, 16

values = torch.randn(batch_size, 1, height, width)
rewards = torch.randn(batch_size, 1, height, width)

# Attention computes similarity-weighted value aggregation
next_values = attn_vi(values, rewards)
print(f"Attention VI output shape: {next_values.shape}")
```

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

### Initialize and Train ScalableVIN Model

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

The ScalableVIN class integrates state encoding, planning, and action prediction. It supports gradient checkpointing for memory efficiency when using high depth values.

```python
import torch
from scaling_vin_pytorch import ScalableVIN

# Initialize the scalable VIN model
scalable_vin = ScalableVIN(
    state_dim=3,              # Number of state channels (e.g., RGB image)
    reward_dim=2,             # Number of reward channels
    num_actions=10,           # Number of possible actions
    num_plans=1,              # Number of parallel value maps
    dim_hidden=150,           # Hidden dimension for reward mapper
    depth=100,                # Number of value iteration layers (can scale to 5000)
    loss_every_num_layers=4,  # Compute loss every N layers for multi-layer supervision
    checkpoint=True,          # Enable gradient checkpointing for memory efficiency
    vi_module_type='dynamic', # Options: 'invariant', 'dynamic', 'attention'
    soft_maxpool=False,       # Use soft max-pool for differentiable action selection
    soft_maxpool_temperature=1.0
)

# Prepare input data
batch_size = 2
grid_size = 32

state = torch.randn(batch_size, 3, grid_size, grid_size)   # State representation
reward = torch.randn(batch_size, 2, grid_size, grid_size)  # Reward map
agent_positions = torch.randint(0, grid_size, (batch_size, 2))  # (y, x) positions
target_actions = torch.randint(0, 10, (batch_size,))  # Ground truth actions

# Training mode: compute loss
loss = scalable_vin(state, reward, agent_positions, target_actions)
loss.backward()
print(f"Training loss: {loss.item():.4f}")

# Inference mode: get action logits
action_logits = scalable_vin(state, reward, agent_positions)
# Shape: (batch_size, depth // loss_every_num_layers, num_actions)
predicted_actions = action_logits[:, 0].argmax(dim=-1)  # Use deepest layer
print(f"Predicted actions: {predicted_actions}")
```

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

### ScalableVIN - Main Planning Model

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

The primary interface for creating scalable value iteration networks that combine state/reward encoding, deep value iteration planning, and action prediction with multi-layer loss computation.

```APIDOC
## ScalableVIN - Main Planning Model

### Description
The primary interface for creating scalable value iteration networks that combine state/reward encoding, deep value iteration planning, and action prediction with multi-layer loss computation.

### Method
POST (or similar, as it's a model initialization and forward pass)

### Endpoint
`/api/scalable_vin` (Hypothetical endpoint for model usage)

### Parameters
#### Initialization Parameters
- **state_dim** (int) - Number of state channels (e.g., RGB image)
- **reward_dim** (int) - Number of reward channels
- **num_actions** (int) - Number of possible actions
- **num_plans** (int) - Number of parallel value maps
- **dim_hidden** (int) - Hidden dimension for reward mapper
- **depth** (int) - Number of value iteration layers (can scale to 5000)
- **loss_every_num_layers** (int) - Compute loss every N layers for multi-layer supervision
- **checkpoint** (bool) - Enable gradient checkpointing for memory efficiency
- **vi_module_type** (str) - Options: 'invariant', 'dynamic', 'attention'
- **soft_maxpool** (bool) - Use soft max-pool for differentiable action selection
- **soft_maxpool_temperature** (float) - Temperature for soft max-pooling

#### Input Data (for forward pass)
- **state** (torch.Tensor) - State representation (batch_size, 3, grid_size, grid_size)
- **reward** (torch.Tensor) - Reward map (batch_size, 2, grid_size, grid_size)
- **agent_positions** (torch.Tensor) - (y, x) positions of the agent (batch_size, 2)
- **target_actions** (torch.Tensor, optional) - Ground truth actions for training (batch_size,)

### Request Example (Initialization)
```python
import torch
from scaling_vin_pytorch import ScalableVIN

scalable_vin = ScalableVIN(
    state_dim=3,
    reward_dim=2,
    num_actions=10,
    num_plans=1,
    dim_hidden=150,
    depth=100,
    loss_every_num_layers=4,
    checkpoint=True,
    vi_module_type='dynamic',
    soft_maxpool=False,
    soft_maxpool_temperature=1.0
)
```

### Request Example (Training Forward Pass)
```python
# Prepare input data
batch_size = 2
grid_size = 32
state = torch.randn(batch_size, 3, grid_size, grid_size)
reward = torch.randn(batch_size, 2, grid_size, grid_size)
agent_positions = torch.randint(0, grid_size, (batch_size, 2))
target_actions = torch.randint(0, 10, (batch_size,))

# Training mode: compute loss
loss = scalable_vin(state, reward, agent_positions, target_actions)
loss.backward()
print(f"Training loss: {loss.item():.4f}")
```

### Request Example (Inference Forward Pass)
```python
# Inference mode: get action logits
action_logits = scalable_vin(state, reward, agent_positions)
```

### Response
#### Success Response (200)
- **loss** (torch.Tensor) - Computed loss value during training.
- **action_logits** (torch.Tensor) - Logits for predicted actions during inference. Shape: (batch_size, depth // loss_every_num_layers, num_actions)

#### Response Example (Training)
```json
{
  "loss": 0.1234
}
```

#### Response Example (Inference)
```json
{
  "action_logits": [[[-0.1, 0.5, ...], ...]]
}
```
```

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

### ActionSelector for Q-Value Based Action Selection

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

Selects actions from Q-value tensors using either hard max or soft max-pooling. Configure `soft_maxpool` and `soft_maxpool_temperature` for desired selection behavior.

```python
import torch
from scaling_vin_pytorch import ActionSelector

# Create action selector with soft max-pooling
selector = ActionSelector(
    num_plans=2,
    soft_maxpool=True,
    soft_maxpool_temperature=0.5  # Lower temp = sharper selection
)

batch_size = 4
num_plans = 2
num_actions = 8
height, width = 16, 16

# Q-values with shape (batch, plans * actions, height, width)
q_values = torch.randn(batch_size, num_plans * num_actions, height, width)

# Select best action per plan
selected_values = selector(q_values)
# Output shape: (batch, num_plans, height, width)
print(f"Q-values shape: {q_values.shape}")
print(f"Selected values shape: {selected_values.shape}")

# With hard max (default)
hard_selector = ActionSelector(num_plans=2, soft_maxpool=False)
hard_selected = hard_selector(q_values)
```

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

### Citations for Scaling Value Iteration Networks

Source: https://github.com/lucidrains/scaling-vin-pytorch/blob/main/README.md

BibTeX references for the Scaling Value Iteration Networks paper and related works.

```bibtex
@article{Wang2024ScalingVI,
    title   = {Scaling Value Iteration Networks to 5000 Layers for Extreme Long-Term Planning},
    author  = {Yuhui Wang and Qingyuan Wu and Weida Li and Dylan R. Ashley and Francesco Faccio and Chao Huang and J{"u}rgen Schmidhuber},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2406.08404},
    url     = {https://api.semanticscholar.org/CorpusID:270391752}
}
```

```bibtex
@misc{pflueger2018soft,
    title   = {Soft Value Iteration Networks for Planetary Rover Path Planning},
    author  = {Max Pflueger and Ali Agha and Gaurav S. Sukhatme},
    year    = {2018},
    url     = {https://openreview.net/forum?id=Sktm4zWRb},
}
```

```bibtex
@inproceedings{Tamar2016ValueIN,
    title   = {Value Iteration Networks},
    author  = {Aviv Tamar and Sergey Levine and P. Abbeel and Yi Wu and Garrett Thomas},
    booktitle = {Neural Information Processing Systems},
    year    = {2016},
    url     = {https://api.semanticscholar.org/CorpusID:11374605}
}
```

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

### ValueIterationNetwork for Deep Recurrent Planning

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

Wraps a ValueIteration module to apply it recurrently for a specified depth, with optional gradient checkpointing for memory efficiency. Use for deep planning in complex environments.

```python
import torch
from scaling_vin_pytorch import ValueIterationNetwork, ValueIteration

# Create the base VI module
vi_module = ValueIteration(
    action_channels=4,
    num_plans=1,
    receptive_field=3
)

# Wrap it in a recurrent network
vin = ValueIterationNetwork(
    vi_module=vi_module,
    depth=50,                   # Number of recurrent iterations
    checkpoint=True,            # Enable gradient checkpointing
    checkpoint_segments=10      # Number of checkpoint segments
)

batch_size = 2
height, width = 32, 32

values = torch.randn(batch_size, 1, height, width)
rewards = torch.randn(batch_size, 1, height, width)

# Run value iteration for specified depth
all_layer_values = vin(values, rewards, depth=50)
# Returns list of value tensors, one per layer
print(f"Number of layer outputs: {len(all_layer_values)}")
print(f"Each layer shape: {all_layer_values[0].shape}")
```

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

### Implement Dynamic Value Iteration

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

DynamicValueIteration computes data-dependent kernels internally, allowing the transition model to adapt based on local state information.

```python
import torch
from scaling_vin_pytorch import DynamicValueIteration

# Create dynamic value iteration module
dynamic_vi = DynamicValueIteration(
    action_channels=8,
    num_plans=2,                # Use multiple parallel plans
    receptive_field=3,
    pad_value=0.0,
    soft_maxpool=True,          # Enable soft max-pooling
    soft_maxpool_temperature=0.5
)

# Input tensors with multiple plans
batch_size = 4
height, width = 16, 16
num_plans = 2

values = torch.randn(batch_size, num_plans, height, width)
rewards = torch.randn(batch_size, num_plans, height, width)

# Dynamic transition computes data-dependent kernels internally
next_values = dynamic_vi(values, rewards)
print(f"Dynamic VI output shape: {next_values.shape}")
```

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

### Implement Invariant Value Iteration

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

ValueIteration uses fixed convolutional kernels for environments with consistent transition dynamics. The receptive field must be an odd integer.

```python
import torch
from scaling_vin_pytorch import ValueIteration

# Create invariant value iteration module
vi_module = ValueIteration(
    action_channels=8,          # Number of action channels
    num_plans=1,                # Number of parallel value plans
    receptive_field=3,          # Kernel size for transitions (must be odd)
    pad_value=0.0,              # Padding value for boundaries
    soft_maxpool=False,         # Use hard max for action selection
    soft_maxpool_temperature=1.0
)

# Input tensors
batch_size = 4
height, width = 16, 16
num_plans = 1

values = torch.randn(batch_size, num_plans, height, width)  # Current value estimates
rewards = torch.randn(batch_size, num_plans, height, width) # Reward signal

# Perform one step of value iteration
next_values = vi_module(values, rewards)
# Output shape: (batch_size, num_plans, height, width)
print(f"Input values shape: {values.shape}")
print(f"Output values shape: {next_values.shape}")
```

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

### DynamicValueIteration - Data-Dependent Transitions

Source: https://context7.com/lucidrains/scaling-vin-pytorch/llms.txt

Dynamic value iteration using input-dependent convolutional kernels, enabling adaptive transition models that change based on local state information.

```APIDOC
## DynamicValueIteration - Data-Dependent Transitions

### Description
Dynamic value iteration using input-dependent convolutional kernels, enabling adaptive transition models that change based on local state information.

### Method
POST (or similar, as it's a module initialization and forward pass)

### Endpoint
`/api/value_iteration/dynamic` (Hypothetical endpoint for module usage)

### Parameters
#### Initialization Parameters
- **action_channels** (int) - Number of action channels
- **num_plans** (int) - Number of parallel value plans
- **receptive_field** (int) - Kernel size for transitions (must be odd)
- **pad_value** (float) - Padding value for boundaries
- **soft_maxpool** (bool) - Enable soft max-pooling
- **soft_maxpool_temperature** (float) - Temperature for soft max-pooling

#### Input Data (for forward pass)
- **values** (torch.Tensor) - Current value estimates (batch_size, num_plans, height, width)
- **rewards** (torch.Tensor) - Reward signal (batch_size, num_plans, height, width)

### Request Example (Initialization)
```python
import torch
from scaling_vin_pytorch import DynamicValueIteration

dynamic_vi = DynamicValueIteration(
    action_channels=8,
    num_plans=2,
    receptive_field=3,
    pad_value=0.0,
    soft_maxpool=True,
    soft_maxpool_temperature=0.5
)
```

### Request Example (Forward Pass)
```python
# Input tensors with multiple plans
batch_size = 4
height, width = 16, 16
num_plans = 2

values = torch.randn(batch_size, num_plans, height, width)
rewards = torch.randn(batch_size, num_plans, height, width)

# Dynamic transition computes data-dependent kernels internally
next_values = dynamic_vi(values, rewards)
```

### Response
#### Success Response (200)
- **next_values** (torch.Tensor) - Updated value estimates after one iteration. Shape: (batch_size, num_plans, height, width)

#### Response Example
```json
{
  "next_values": [[[...]]]
}
```
```

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