LeRobot Notes: The Basic Workflow

1
dataset = LeRobotDataset("lerobot/pusht")
2
print(dataset)
3

4
# output
5
# LeRobotDataset({
6
#     Repository ID: 'lerobot/pusht',
7
#     Number of selected episodes: '206',
8
#     Number of selected samples: '25650',
9
#     Features: '['observation.image', 'observation.state', 'action', 'episode_index', 'frame_index', 'timestamp', 'next.reward', 'next.done', 'next.success', 'index', 'task_index']',
10
# })',

1
python src/lerobot/scripts/lerobot_dataset_viz.py \
2
    --repo-id lerobot/pusht \
3
    --episode-index 1 \
4
    --display-compressed-images true

1
dataset = StreamingLeRobotDataset("lerobot/pusht")
2
for sample in dataset:
3
    print(sample['action'])

1
def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
2
    # ...
3

4
    # buffer 内打散
5
    buffer_indices_generator = self._iter_random_indices(rng, self.buffer_size)
6

7
    # 拆成多个 shard, backtrack_dataset 是和 delta_timestamp 相关，需要拿到前后的若干条数据
8
    idx_to_backtrack_dataset = {
9
        idx: self._make_backtrackable_dataset(safe_shard(self.hf_dataset, idx, self.num_shards))
10
        for idx in range(self.num_shards)
11
    }
12

13
    frames_buffer = []
14
    while available_shards := list(idx_to_backtrack_dataset.keys()):
15
        # 随便选一个 shard，拿到对应的 backtrack_dataset
16
        shard_key = next(self._infinite_generator_over_elements(rng, available_shards))
17
        backtrack_dataset = idx_to_backtrack_dataset[shard_key]  # selects which shard to iterate on
18

19
        try:
20
            # make_frame 有可能会产生多个 frame，和 delta_timestamp 有关，下面再详细说说
21
            for frame in self.make_frame(backtrack_dataset):
22
                # 蓄水池采样，如果满了就丢出去第 i 个，把当前的 frame 补到 i-th 里
23
                if len(frames_buffer) == self.buffer_size:
24
                    i = next(buffer_indices_generator)  # samples a element from the buffer
25
                    yield frames_buffer[i]
26
                    frames_buffer[i] = frame
27
                else:
28
                    frames_buffer.append(frame)
29
                break  # random shard sampled, switch shard
30
        except (
31
            RuntimeError,
32
            StopIteration,
33
        ):  # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7
34
            del idx_to_backtrack_dataset[shard_key]  # Remove exhausted shard, onto another shard

1
self._source: Iterator[T] = iter(iterable)
2
self._back_buf: deque[T] = deque(maxlen=history)
3
self._ahead_buf: deque[T] = deque(maxlen=lookahead) if lookahead > 0 else deque()
4
self._cursor: int = 0

1
for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=True):
2
    loaded_frames.append(frame)
3
    loaded_ts.append(pts.item())
4
    if log_loaded_timestamps:
5
        logging.info(f"Frame loaded at timestamp={pts:.4f}")
6

7
query_ts = torch.tensor(timestamps)
8
loaded_ts = torch.tensor(loaded_ts)
9

10
# compute distances between each query timestamp and loaded timestamps
11
dist = torch.cdist(query_ts[:, None], loaded_ts[:, None], p=1)
12
min_, argmin_ = dist.min(1)

1
# preprocessor 和 postprocessor
2
preprocessor, postprocessor = make_pre_post_processors(
3
    policy_cfg=cfg.policy,
4
    pretrained_path=cfg.policy.pretrained_path,
5
    **processor_kwargs,
6
    **postprocessor_kwargs,
7
)
8
logging.info(f"preprocessor: {preprocessor}, postprocessor: {postprocessor}")
9

10
# ...
11

12
# accelerator 根据加速器（如 deepspeed）等的封装，初始化几个关键对象
13
policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
14
    policy, optimizer, dataloader, lr_scheduler
15
)
16
print(f"policy: {policy}, optimizer: {optimizer}, dataloader: {dataloader}, lr_scheduler: {lr_scheduler}")
17

18
# ...
19

20
for _ in range(step, cfg.steps):
21
    start_time = time.perf_counter()
22
    # dl_iter 是上面讲到的 dataloader
23
    batch = next(dl_iter)
24

25
    batch = preprocessor(batch)
26
    train_tracker.dataloading_s = time.perf_counter() - start_time
27

28
    train_tracker, output_dict = update_policy(
29
        train_tracker,
30
        policy,
31
        batch,
32
        optimizer,
33
        cfg.optimizer.grad_clip_norm,
34
        accelerator=accelerator,
35
        lr_scheduler=lr_scheduler,
36
        rabc_weights_provider=rabc_weights,
37
    )

1
input_steps = [
2
    RenameObservationsProcessorStep(rename_map={}),
3
    AddBatchDimensionProcessorStep(),
4
    DeviceProcessorStep(device=config.device),
5
    NormalizerProcessorStep(
6
        features={**config.input_features, **config.output_features},
7
        norm_map=config.normalization_mapping,
8
        stats=dataset_stats,
9
        device=config.device,
10
    ),
11
]
12
output_steps = [
13
    UnnormalizerProcessorStep(
14
        features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
15
    ),
16
    DeviceProcessorStep(device="cpu"),
17
]
18

19
return (
20
    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
21
        steps=input_steps,
22
        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
23
    ),
24
    PolicyProcessorPipeline[PolicyAction, PolicyAction](
25
        steps=output_steps,
26
        name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
27
        to_transition=policy_action_to_transition,
28
        to_output=transition_to_policy_action,
29
    ),
30
)

1
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, tuple[Tensor, Tensor] | tuple[None, None]]:
2
    """A forward pass through the Action Chunking Transformer (with optional VAE encoder).
3

4
    `batch` should have the following structure:
5
    {
6
        [robot_state_feature] (optional): (B, state_dim) batch of robot states.
7

8
        [image_features]: (B, n_cameras, C, H, W) batch of images.
9
            AND/OR
10
        [env_state_feature]: (B, env_dim) batch of environment states.
11

12
        [action_feature] (optional, only if training with VAE): (B, chunk_size, action dim) batch of actions.
13
    }
14

15
    Returns:
16
        (B, chunk_size, action_dim) batch of action sequences
17
        Tuple containing the latent PDF's parameters (mean, log(σ²)) both as (B, L) tensors where L is the
18
        latent dimension.
19
    """

1
decoder_in = torch.zeros(
2
    (self.config.chunk_size, batch_size, self.config.dim_model),
3
    dtype=encoder_in_pos_embed.dtype,
4
    device=encoder_in_pos_embed.device,
5
)
6
decoder_out = self.decoder(
7
    decoder_in,
8
    encoder_out,
9
    encoder_pos_embed=encoder_in_pos_embed,
10
    decoder_pos_embed=self.decoder_pos_embed.weight.unsqueeze(1),
11
)
12

13
#...
14

15
# 最后用 action_head 把 model output 转成有限维度的 action
16
# Final action regression head on the output of the transformer's decoder.
17
self.action_head = nn.Linear(config.dim_model, self.config.action_feature.shape[0])