### Install Q-Transformer

Source: https://github.com/lucidrains/q-transformer/blob/main/README.md

Install the package via pip.

```bash
$ pip install q-transformer
```

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

### Setup Mock Environment

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Initializes a mock environment for testing purposes, defining the shapes for state and text embeddings. Requires Q-Transformer library.

```python
# 2. Setup environment
env = MockEnvironment(
    state_shape=(3, 6, 224, 224),
    text_embed_shape=(768,)
)
```

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

### Initialize and Use QRoboticTransformer Model

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Initializes the QRoboticTransformer model with specified configurations for vision backbone and transformer layers. Demonstrates how to get optimal actions autoregressively and calculate Q-values for given actions.

```python
import torch
from q_transformer import QRoboticTransformer

# Initialize the Q-Transformer model
model = QRoboticTransformer(
    vit=dict(
        num_classes=1000,
        dim_conv_stem=64,
        dim=64,
        dim_head=64,
        depth=(2, 2, 5, 2),
        window_size=7,
        mbconv_expansion_rate=4,
        mbconv_shrinkage_rate=0.25,
        dropout=0.1
    ),
    num_actions=8,           # Number of action dimensions (e.g., 8 for robot arm joints)
    action_bins=256,         # Discretization bins per action
    depth=6,                 # Transformer depth
    heads=8,                 # Attention heads
    dim_head=64,             # Dimension per head
    cond_drop_prob=0.2,      # Classifier-free guidance dropout
    dueling=True,            # Use dueling Q-network architecture
    condition_on_text=True,  # Enable text conditioning
    dual_critics=False,      # Double Q-learning for reduced overestimation
    flash_attn=True          # Use Flash Attention
)

# Input: batch of videos (batch, channels, frames, height, width)
video = torch.randn(2, 3, 6, 224, 224)

# Text instructions for conditioning
instructions = [
    'bring me that apple sitting on the table',
    'please pass the butter'
]

# Get optimal actions autoregressively
actions = model.get_optimal_actions(video, instructions)
# Output shape: (batch_size, num_actions) -> tensor of action bin indices

# Get Q-values for specific actions
q_values = model(video, instructions, actions=actions)
# Output shape: (batch_size, num_actions, action_bins)
```

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

### Get Actions with Epsilon-Greedy Exploration

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Provides a wrapper for action selection in QRoboticTransformer that includes epsilon-greedy exploration. Allows for a specified probability of selecting random actions instead of optimal ones.

```python
import torch
from q_transformer import QRoboticTransformer

model = QRoboticTransformer(
    vit=dict(num_classes=1000, dim_conv_stem=64, dim=64, dim_head=64,
        depth=(2, 2, 5, 2), window_size=7, mbconv_expansion_rate=4,
        mbconv_shrinkage_rate=0.25, dropout=0.1
    ),
    num_actions=8,
    action_bins=256,
    depth=1,
    heads=8,
    dim_head=64
)

video = torch.randn(2, 3, 6, 224, 224)
instructions = ['move forward', 'turn left']

# Pure exploitation (no exploration)
actions = model.get_actions(video, instructions, prob_random_action=0.0)

# Epsilon-greedy exploration (25% random actions)
actions = model.get_actions(video, instructions, prob_random_action=0.25)
```

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

### Get Optimal Actions with QRoboticTransformer

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Demonstrates how to autoregressively decode optimal actions using the QRoboticTransformer model. Shows basic action selection, returning Q-values, and conditioning on partial actions for hierarchical control.

```python
import torch
from q_transformer import QRoboticTransformer

model = QRoboticTransformer(
    vit=dict(num_classes=1000, dim_conv_stem=64, dim=64, dim_head=64,
        depth=(2, 2, 5, 2), window_size=7, mbconv_expansion_rate=4,
        mbconv_shrinkage_rate=0.25, dropout=0.1
    ),
    num_actions=8,
    action_bins=256,
    depth=1,
    heads=8,
    dim_head=64
)

video = torch.randn(4, 3, 6, 224, 224)
instructions = ['pick up the red block'] * 4

# Basic optimal action selection
optimal_actions = model.get_optimal_actions(video, instructions)
# Shape: (4, 8) - 4 batches, 8 action dimensions

# Get actions with Q-values for debugging/analysis
actions, q_values = model.get_optimal_actions(
    video,
    instructions,
    return_q_values=True
)
# actions shape: (4, 8)
# q_values shape: (4, 8, 256) - Q-values for each action bin

# Continue from partial actions (useful for hierarchical control)
partial_actions = actions[:, :3]  # First 3 actions
remaining_actions = model.get_optimal_actions(
    video,
    instructions,
    actions=partial_actions  # Condition on previous actions
)
```

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

### Inference with Trained Q-Transformer Model

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Deploys the trained model for inference to get optimal actions based on video input and textual instructions. Uses PyTorch for tensor operations and requires the model to be in evaluation mode.

```python
# 5. Deploy trained model for inference
model.eval()
video = torch.randn(1, 3, 6, 224, 224)
instruction = ['pick up the red cube']

with torch.no_grad():
    actions = model.get_optimal_actions(video, instruction)
    # actions: tensor of shape (1, 8) with action bin indices

    # Convert to continuous actions if needed
    continuous_actions = (actions.float() / 256.0) * 2 - 1  # Scale to [-1, 1]
```

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

### Initialize QRoboticTransformer Model

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Sets up the QRoboticTransformer model with specified parameters for vision and action components. Requires PyTorch and Q-Transformer library.

```python
import torch
from q_transformer import (
    QRoboticTransformer,
    QLearner,
    Agent,
    ReplayMemoryDataset
)
from q_transformer.mocks import MockEnvironment

# 1. Initialize model
model = QRoboticTransformer(
    vit=dict(
        num_classes=1000,
        dim_conv_stem=64,
        dim=64,
        dim_head=64,
        depth=(2, 2, 5, 2),
        window_size=7,
        mbconv_expansion_rate=4,
        mbconv_shrinkage_rate=0.25,
        dropout=0.1
    ),
    num_actions=8,
    action_bins=256,
    depth=1,
    heads=8,
    dim_head=64,
    cond_drop_prob=0.2,
    dueling=True
)
```

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

### Initialize and Train Q-Transformer

Source: https://github.com/lucidrains/q-transformer/blob/main/README.md

Configure the QRoboticTransformer model, interact with an environment using the Agent class, and perform training with QLearner.

```python
import torch

from q_transformer import (
	QRoboticTransformer,
	QLearner,
	Agent,
	ReplayMemoryDataset
)

# the attention model

model = QRoboticTransformer(
	vit = dict(
		num_classes = 1000,
		dim_conv_stem = 64,
		dim = 64,
		dim_head = 64,
		depth = (2, 2, 5, 2),
		window_size = 7,
		mbconv_expansion_rate = 4,
		mbconv_shrinkage_rate = 0.25,
		dropout = 0.1
	),
	num_actions = 8,
	action_bins = 256,
	depth = 1,
	heads = 8,
	dim_head = 64,
	cond_drop_prob = 0.2,
	dueling = True
)

# you need to supply your own environment, by overriding BaseEnvironment

from q_transformer.mocks import MockEnvironment

env = MockEnvironment(
	state_shape = (3, 6, 224, 224),
	text_embed_shape = (768,)
)

# env.init()     should return instructions and initial state: Tuple[str, Tensor[*state_shape]]
# env(actions)   should return rewards, next state, and done flag: Tuple[Tensor[()], Tensor[*state_shape], Tensor[()]]

# agent is a class that allows the q-model to interact with the environment to generate a replay memory dataset for learning

agent = Agent(
	model,
	environment = env,
	num_episodes = 1000,
	max_num_steps_per_episode = 100,
)

agent()

# Q learning on the replay memory dataset on the model

q_learner = QLearner(
	model,
	dataset = ReplayMemoryDataset(),
	num_train_steps = 10000,
	learning_rate = 3e-4,
	batch_size = 4,
	grad_accum_every = 16,
)

q_learner()

# after much learning
# your robot should be better at selecting optimal actions

video = torch.randn(2, 3, 6, 224, 224)

instructions = [
	'bring me that apple sitting on the table',
	'please pass the butter'
]

actions = model.get_optimal_actions(video, instructions)
```

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

### Train Model with Q-Learning

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Configures and runs the QLearner for training the model using collected experiences from a ReplayMemoryDataset. Requires Q-Transformer library.

```python
# 4. Train with Q-Learning
q_learner = QLearner(
    model,
    dataset=ReplayMemoryDataset(),
    num_train_steps=10000,
    learning_rate=3e-4,
    batch_size=4,
    grad_accum_every=16
)
q_learner()  # Runs training loop
```

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

### Train Model with QLearner

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Configures and runs the training loop using the QLearner wrapper, supporting EMA target networks and checkpointing.

```python
import torch
from q_transformer import QRoboticTransformer, QLearner
from q_transformer.mocks import MockReplayDataset

# Create model
model = QRoboticTransformer(
	vit=dict(
		num_classes=1000, dim_conv_stem=64, dim=64, dim_head=64,
		depth=(2, 2, 5, 2), window_size=7, mbconv_expansion_rate=4,
		mbconv_shrinkage_rate=0.25, dropout=0.1
	),
	num_actions=8,
	action_bins=256,
	depth=1,
	heads=8,
	dim_head=64,
	dueling=True
)

# Create Q-Learner with dataset
q_learner = QLearner(
	model,
	dataset=MockReplayDataset(
		length=10000,
		num_actions=8,
		num_action_bins=256
	),
	num_train_steps=10000,
	learning_rate=3e-4,
	batch_size=4,
	grad_accum_every=16,
	discount_factor_gamma=0.98,
	conservative_reg_loss_weight=1.0,  # CQL regularization
	n_step_q_learning=True,            # Enable n-step returns
	use_bce_loss=True,                 # Classification loss (from "Stop Regressing" paper)
	checkpoint_folder='./checkpoints',
	checkpoint_every=1000,
	q_target_ema_kwargs=dict(
		beta=0.99,
		update_after_step=10,
		update_every=5
	)
)

# Run training loop
q_learner()

# Save and load checkpoints
q_learner.save(checkpoint_num=1)
q_learner.load('./checkpoints/checkpoint-1.pt')
```

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

### get_actions

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Wrapper for action selection with epsilon-greedy exploration.

```APIDOC
## get_actions

### Description
Selects actions using epsilon-greedy exploration. Returns random actions with probability `prob_random_action`.

### Parameters
- **video** (torch.Tensor) - Batch of videos.
- **instructions** (list) - Text instructions.
- **prob_random_action** (float) - Probability of selecting a random action.
```

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

### Define Custom Environment with BaseEnvironment

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Implement the BaseEnvironment class by defining init() and forward() methods to handle environment state and action execution.

```python
import torch
from torch import Tensor
from q_transformer.agent import BaseEnvironment
from q_transformer.tensor_typing import Float, Bool

class RobotArmEnvironment(BaseEnvironment):
    def __init__(self):
        super().__init__(
            state_shape=(3, 6, 224, 224),  # RGB video: channels, frames, H, W
            text_embed_shape=(768,)         # Text embedding size
        )
        self.goal_position = None
        self.current_position = None

    def init(self) -> tuple[str, Tensor]:
        """Reset environment and return (instruction, initial_state)"""
        self.goal_position = torch.rand(3)
        self.current_position = torch.zeros(3)

        instruction = "move the robot arm to the target position"
        initial_state = self._get_camera_observation()
        return instruction, initial_state

    def forward(self, actions: Tensor) -> tuple[Float[''], Float['...'], Bool['']]:
        """
        Execute actions and return (reward, next_state, done).

        Args:
            actions: Tensor of action bin indices, shape (num_actions,)

        Returns:
            reward: Scalar reward tensor
            next_state: Next observation tensor
            done: Boolean indicating episode termination
        """
        # Convert discrete actions to continuous (example)
        continuous_actions = (actions.float() / 256.0) * 2 - 1  # Scale to [-1, 1]

        # Update robot state
        self.current_position += continuous_actions[:3] * 0.1

        # Calculate reward (negative distance to goal)
        distance = torch.norm(self.current_position - self.goal_position)
        reward = -distance

        # Check if done
        done = distance < 0.1

        # Get next observation
        next_state = self._get_camera_observation()

        return reward, next_state, done

    def _get_camera_observation(self) -> Tensor:
        # Return camera frames (placeholder)
        return torch.randn(self.state_shape, device=self.device)

# Usage
env = RobotArmEnvironment()
instruction, state = env.init()
# instruction: "move the robot arm to the target position"
# state shape: (3, 6, 224, 224)
```

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

### Generate Random Actions with Q-Transformer

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Demonstrates how to generate random actions for exploration using the model's internal methods.

```python
# Pure exploration (all random actions)
random_actions = model.get_actions(video, instructions, prob_random_action=1.0)

# Get random actions directly
random_actions = model.get_random_actions(batch_size=2)
# Shape: (2, 8) - random action bins for each action dimension
```

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

### Initialize and Use MaxViT Vision Backbone

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Configure the MaxViT vision encoder for image processing or text-conditioned tasks.

```python
import torch
from q_transformer.q_robotic_transformer import MaxViT

# Create standalone MaxViT for vision encoding
vit = MaxViT(
    num_classes=1000,
    dim=64,
    dim_conv_stem=64,
    dim_head=64,
    depth=(2, 2, 5, 2),      # Blocks per stage
    window_size=7,            # Attention window size
    mbconv_expansion_rate=4,
    mbconv_shrinkage_rate=0.25,
    dropout=0.1,
    channels=3,
    flash_attn=True,
    use_layernorm=True        # Use LayerNorm instead of BatchNorm
)

# Process images
images = torch.randn(4, 3, 224, 224)

# Get classification logits
logits = vit(images)
# Shape: (4, 1000)

# Get embeddings for downstream tasks
embeddings = vit(images, return_embeddings=True)
# Shape: (4, embed_dim, H', W') where H', W' are downsampled

# With text conditioning (for QRoboticTransformer integration)
texts = ['robot arm', 'gripper', 'table', 'cube']
logits_conditioned = vit(images, texts=texts, cond_drop_prob=0.1)
```

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

### Collect Experiences with Agent

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Uses the Agent class to interact with the environment, collect experiences, and store them in replay memories. Requires the initialized model and environment.

```python
# 3. Collect experiences with Agent
agent = Agent(
    model,
    environment=env,
    num_episodes=1000,
    max_num_steps_per_episode=100
)
agent()  # Stores memories to ./replay_memories_data/
```

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

### QRoboticTransformer Initialization

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Initializes the main Q-Transformer model, combining a MaxViT vision encoder with a transformer decoder for Q-value prediction.

```APIDOC
## QRoboticTransformer

### Description
Initializes the QRoboticTransformer model architecture.

### Parameters
- **vit** (dict) - Configuration for the MaxViT vision backbone.
- **num_actions** (int) - Number of action dimensions.
- **action_bins** (int) - Discretization bins per action.
- **depth** (int) - Transformer depth.
- **heads** (int) - Number of attention heads.
- **dim_head** (int) - Dimension per head.
- **cond_drop_prob** (float) - Classifier-free guidance dropout probability.
- **dueling** (bool) - Whether to use dueling Q-network architecture.
- **condition_on_text** (bool) - Enable text conditioning.
- **dual_critics** (bool) - Use double Q-learning.
- **flash_attn** (bool) - Use Flash Attention.
```

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

### Load Replay Memory with ReplayMemoryDataset

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Use ReplayMemoryDataset to load data from memory-mapped files for training, supporting both single-step and n-step returns.

```python
from q_transformer import ReplayMemoryDataset
from torch.utils.data import DataLoader

# Load replay memories (after Agent has collected data)
dataset = ReplayMemoryDataset(
    folder='./replay_memories_data',
    num_timesteps=1  # Single-step Q-learning
)

# For n-step Q-learning
n_step_dataset = ReplayMemoryDataset(
    folder='./replay_memories_data',
    num_timesteps=4  # 4-step returns
)

# Create dataloader for training
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

# Each batch contains:
# - text_embeds: (batch, timesteps, embed_dim)
# - states: (batch, timesteps, channels, frames, height, width)
# - actions: (batch, timesteps, num_actions)
# - next_state: (batch, channels, frames, height, width)
# - next_text_embed: (batch, embed_dim)
# - rewards: (batch, timesteps)
# - dones: (batch, timesteps)

for batch in dataloader:
    text_embeds, states, actions, next_state, next_text_embed, rewards, dones = batch
    # Process batch for training
    break
```

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

### Collect Experiences with Agent

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Uses the Agent class to interact with an environment and store experiences in memory-mapped files.

```python
import torch
from q_transformer import QRoboticTransformer, Agent
from q_transformer.mocks import MockEnvironment

# Create model
model = QRoboticTransformer(
	vit=dict(
		num_classes=1000, dim_conv_stem=64, dim=64, dim_head=64,
		depth=(2, 2, 5, 2), window_size=7, mbconv_expansion_rate=4,
		mbconv_shrinkage_rate=0.25, dropout=0.1
	),
	num_actions=8,
	action_bins=256,
	depth=1,
	heads=8,
	dim_head=64
)

# Create environment (extend BaseEnvironment for custom envs)
env = MockEnvironment(
	state_shape=(3, 6, 224, 224),    # (channels, frames, height, width)
	text_embed_shape=(768,)           # Text embedding dimension
)

# Create agent for data collection
agent = Agent(
	model,
	environment=env,
	memories_dataset_folder='./replay_memories_data',
	num_episodes=1000,
	max_num_steps_per_episode=100,
	epsilon_start=0.25,      # Initial exploration rate
	epsilon_end=0.001,       # Final exploration rate
	num_steps_to_target_epsilon=1000
)

# Run data collection (stores to memory-mapped files)
agent()
# Output: replay memories stored to ./replay_memories_data/
```

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

### get_optimal_actions

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Autoregressively decodes optimal actions by selecting the action with the maximum Q-value at each step.

```APIDOC
## get_optimal_actions

### Description
Decodes optimal actions autoregressively. Supports returning Q-values and conditioning on partial actions.

### Parameters
- **video** (torch.Tensor) - Batch of videos (batch, channels, frames, height, width).
- **instructions** (list) - Text instructions for conditioning.
- **return_q_values** (bool) - Whether to return Q-values alongside actions.
- **actions** (torch.Tensor) - Optional partial actions to condition on.
```

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

### Pre-compute Text Embeddings for Inference

Source: https://context7.com/lucidrains/q-transformer/llms.txt

Caches text embeddings to speed up inference and forward passes when instructions are reused.

```python
import torch
from q_transformer import QRoboticTransformer

model = QRoboticTransformer(
	vit=dict(
		num_classes=1000, dim_conv_stem=64, dim=64, dim_head=64,
		depth=(2, 2, 5, 2), window_size=7, mbconv_expansion_rate=4,
		mbconv_shrinkage_rate=0.25, dropout=0.1
	),
	num_actions=8,
	action_bins=256,
	depth=1,
	heads=8,
	dim_head=64,
	condition_on_text=True
)

# Pre-compute text embeddings
instructions = [
	'pick up the cup',
	'place it on the table',
	'push the button'
]
text_embeds = model.embed_texts(instructions)
# Shape depends on text conditioner, typically (batch, embed_dim)

# Use cached embeddings for inference (faster than re-encoding)
video = torch.randn(3, 3, 6, 224, 224)
actions = model.get_optimal_actions(video, text_embeds=text_embeds)

# Can also use with forward pass
q_values = model(video, text_embeds=text_embeds, actions=actions)
```

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