my-dataset/
├── data/
│   └── ...
├── meta/
│   ├── info.json
│   ├── stats.json
│   ├── tasks.parquet
│   ├── subtasks.parquet      # Subtask index → subtask string mapping
│   └── episodes/
│       └── ...
└── videos/
    └── ...

# Example frame data in the parquet file
{
    "index": 42,
    "timestamp": 1.4,
    "episode_index": 0,
    "task_index": 0,
    "subtask_index": 2,  # References "Lift the apple"
    "observation.state": [...],
    "action": [...],
}

from lerobot.datasets.lerobot_dataset import LeRobotDataset

# Load a dataset with subtask annotations
dataset = LeRobotDataset("jadechoghari/collect-fruit-annotated")

# Access a sample
sample = dataset[100]

# The sample includes both task and subtask information
print(sample["task"])        # "Collect the fruit"
print(sample["subtask"])     # "Grasp the apple"
print(sample["task_index"])  # tensor(0)
print(sample["subtask_index"])  # tensor(2)

# Check if subtasks are available
has_subtasks = (
    "subtask_index" in dataset.features
    and dataset.meta.subtasks is not None
)

if has_subtasks:
    print(f"Dataset has {len(dataset.meta.subtasks)} unique subtasks")
    print("Subtasks:", list(dataset.meta.subtasks.index))

from lerobot.processor.tokenizer_processor import TokenizerProcessor
from lerobot.processor.pipeline import ProcessorPipeline

# Create a tokenizer processor
tokenizer_processor = TokenizerProcessor(
    tokenizer_name_or_path="google/paligemma-3b-pt-224",
    padding="max_length",
    max_length=64,
)

# The processor will automatically tokenize subtasks if present in the batch
# and add them to the observation under:
# - "observation.subtask.tokens"
# - "observation.subtask.attention_mask"

import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset

dataset = LeRobotDataset("jadechoghari/collect-fruit-annotated")

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=16,
    shuffle=True,
)

for batch in dataloader:
    # Access subtask information in the batch
    subtasks = batch["subtask"]  # List of subtask strings
    subtask_indices = batch["subtask_index"]  # Tensor of subtask indices

    # Use for training hierarchical policies or reward models
    print(f"Batch subtasks: {set(subtasks)}")

from lerobot.datasets.lerobot_dataset import LeRobotDataset

# Example dataset with subtask annotations
dataset = LeRobotDataset("jadechoghari/collect-fruit-annotated")

# Explore the subtasks
print("Available subtasks:")
for subtask_name in dataset.meta.subtasks.index:
    print(f"  - {subtask_name}")

# Get subtask distribution
subtask_counts = {}
for i in range(len(dataset)):
    sample = dataset[i]
    subtask = sample["subtask"]
    subtask_counts[subtask] = subtask_counts.get(subtask, 0) + 1

print("\nSubtask distribution:")
for subtask, count in sorted(subtask_counts.items(), key=lambda x: -x[1]):
    print(f"  {subtask}: {count} frames")

class HierarchicalPolicy(nn.Module):
    def __init__(self, num_subtasks):
        super().__init__()
        self.action_head = nn.Linear(hidden_dim, action_dim)
        self.subtask_head = nn.Linear(hidden_dim, num_subtasks)

    def forward(self, observations):
        features = self.encoder(observations)
        actions = self.action_head(features)
        subtask_logits = self.subtask_head(features)
        return actions, subtask_logits

# SARM predicts:
# - Stage: Which subtask is being executed (discrete)
# - Progress: How far along the subtask (continuous 0-1)

class SARMRewardModel(nn.Module):
    def forward(self, observations):
        features = self.encoder(observations)
        stage_logits = self.stage_classifier(features)
        progress = self.progress_regressor(features)
        return stage_logits, progress

def visualize_execution(model, observations):
    for t, obs in enumerate(observations):
        action, subtask_logits = model(obs)
        predicted_subtask = subtask_names[subtask_logits.argmax()]
        print(f"t={t}: Executing '{predicted_subtask}'")