diff --git a/Makefile b/Makefile index 2bab6199..2a32048e 100644 --- a/Makefile +++ b/Makefile @@ -146,7 +146,7 @@ test-default-ete-eval: test-act-pusht-tutorial: - cp examples/advanced/train_act_pusht/act_pusht.yaml lerobot/configs/policy/created_by_Makefile.yaml + cp examples/advanced/1_train_act_pusht/act_pusht.yaml lerobot/configs/policy/created_by_Makefile.yaml python lerobot/scripts/train.py \ policy=created_by_Makefile.yaml \ env=pusht \ diff --git a/examples/advanced/calculate_validation_loss.py b/examples/advanced/calculate_validation_loss.py deleted file mode 100644 index 1428014b..00000000 --- a/examples/advanced/calculate_validation_loss.py +++ /dev/null @@ -1,90 +0,0 @@ -"""This script demonstrates how to slice a dataset and calculate the loss on a subset of the data. - -This technique can be useful for debugging and testing purposes, as well as identifying whether a policy -is learning effectively. - -Furthermore, relying on validation loss to evaluate performance is generally not considered a good practice, -especially in the context of imitation learning. The most reliable approach is to evaluate the policy directly -on the target environment, whether that be in simulation or the real world. -""" - -import math -from pathlib import Path - -import torch -from huggingface_hub import snapshot_download - -from lerobot.common.datasets.lerobot_dataset import LeRobotDataset -from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy - -device = torch.device("cuda") - -# Download the diffusion policy for pusht environment -pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht")) -# OR uncomment the following to evaluate a policy from the local outputs/train folder. -# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion") - -policy = DiffusionPolicy.from_pretrained(pretrained_policy_path) -policy.eval() -policy.to(device) - -# Set up the dataset. -delta_timestamps = { - # Load the previous image and state at -0.1 seconds before current frame, - # then load current image and state corresponding to 0.0 second. - "observation.image": [-0.1, 0.0], - "observation.state": [-0.1, 0.0], - # Load the previous action (-0.1), the next action to be executed (0.0), - # and 14 future actions with a 0.1 seconds spacing. All these actions will be - # used to calculate the loss. - "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4], -} - -# Load the last 10% of episodes of the dataset as a validation set. -# - Load full dataset -full_dataset = LeRobotDataset("lerobot/pusht", split="train") -# - Calculate train and val subsets -num_train_episodes = math.floor(full_dataset.num_episodes * 90 / 100) -num_val_episodes = full_dataset.num_episodes - num_train_episodes -print(f"Number of episodes in full dataset: {full_dataset.num_episodes}") -print(f"Number of episodes in training dataset (90% subset): {num_train_episodes}") -print(f"Number of episodes in validation dataset (10% subset): {num_val_episodes}") -# - Get first frame index of the validation set -first_val_frame_index = full_dataset.episode_data_index["from"][num_train_episodes].item() -# - Load frames subset belonging to validation set using the `split` argument. -# It utilizes the `datasets` library's syntax for slicing datasets. -# For more information on the Slice API, please see: -# https://huggingface.co/docs/datasets/v2.19.0/loading#slice-splits -train_dataset = LeRobotDataset( - "lerobot/pusht", split=f"train[:{first_val_frame_index}]", delta_timestamps=delta_timestamps -) -val_dataset = LeRobotDataset( - "lerobot/pusht", split=f"train[{first_val_frame_index}:]", delta_timestamps=delta_timestamps -) -print(f"Number of frames in training dataset (90% subset): {len(train_dataset)}") -print(f"Number of frames in validation dataset (10% subset): {len(val_dataset)}") - -# Create dataloader for evaluation. -val_dataloader = torch.utils.data.DataLoader( - val_dataset, - num_workers=4, - batch_size=64, - shuffle=False, - pin_memory=device != torch.device("cpu"), - drop_last=False, -) - -# Run validation loop. -loss_cumsum = 0 -n_examples_evaluated = 0 -for batch in val_dataloader: - batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} - output_dict = policy.forward(batch) - - loss_cumsum += output_dict["loss"].item() - n_examples_evaluated += batch["index"].shape[0] - -# Calculate the average loss over the validation set. -average_loss = loss_cumsum / n_examples_evaluated - -print(f"Average loss on validation set: {average_loss:.4f}") diff --git a/examples/advanced/train_act_pusht/act_pusht.yaml b/examples/advanced/train_act_pusht/act_pusht.yaml deleted file mode 100644 index 38e542fb..00000000 --- a/examples/advanced/train_act_pusht/act_pusht.yaml +++ /dev/null @@ -1,87 +0,0 @@ -# @package _global_ - -# Change the seed to match what PushT eval uses -# (to avoid evaluating on seeds used for generating the training data). -seed: 100000 -# Change the dataset repository to the PushT one. -dataset_repo_id: lerobot/pusht - -override_dataset_stats: - observation.image: - # stats from imagenet, since we use a pretrained vision model - mean: [[[0.485]], [[0.456]], [[0.406]]] # (c,1,1) - std: [[[0.229]], [[0.224]], [[0.225]]] # (c,1,1) - -training: - offline_steps: 80000 - online_steps: 0 - eval_freq: 10000 - save_freq: 100000 - log_freq: 250 - save_model: true - - batch_size: 8 - lr: 1e-5 - lr_backbone: 1e-5 - weight_decay: 1e-4 - grad_clip_norm: 10 - online_steps_between_rollouts: 1 - - delta_timestamps: - action: "[i / ${fps} for i in range(${policy.chunk_size})]" - -eval: - n_episodes: 50 - batch_size: 50 - -# See `configuration_act.py` for more details. -policy: - name: act - - # Input / output structure. - n_obs_steps: 1 - chunk_size: 100 # chunk_size - n_action_steps: 100 - - input_shapes: - observation.image: [3, 96, 96] - observation.state: ["${env.state_dim}"] - output_shapes: - action: ["${env.action_dim}"] - - # Normalization / Unnormalization - input_normalization_modes: - observation.image: mean_std - # Use min_max normalization just because it's more standard. - observation.state: min_max - output_normalization_modes: - # Use min_max normalization just because it's more standard. - action: min_max - - # Architecture. - # Vision backbone. - vision_backbone: resnet18 - pretrained_backbone_weights: ResNet18_Weights.IMAGENET1K_V1 - replace_final_stride_with_dilation: false - # Transformer layers. - pre_norm: false - dim_model: 512 - n_heads: 8 - dim_feedforward: 3200 - feedforward_activation: relu - n_encoder_layers: 4 - # Note: Although the original ACT implementation has 7 for `n_decoder_layers`, there is a bug in the code - # that means only the first layer is used. Here we match the original implementation by setting this to 1. - # See this issue https://github.com/tonyzhaozh/act/issues/25#issue-2258740521. - n_decoder_layers: 1 - # VAE. - use_vae: true - latent_dim: 32 - n_vae_encoder_layers: 4 - - # Inference. - temporal_ensemble_momentum: null - - # Training and loss computation. - dropout: 0.1 - kl_weight: 10.0 diff --git a/examples/advanced/train_act_pusht/train_act_pusht.md b/examples/advanced/train_act_pusht/train_act_pusht.md deleted file mode 100644 index 5ead1610..00000000 --- a/examples/advanced/train_act_pusht/train_act_pusht.md +++ /dev/null @@ -1,64 +0,0 @@ -In this tutorial we will learn how to adapt a policy configuration to be compatible with a new environment and dataset. As a concrete example, we will adapt the default configuration for ACT to be compatible with the PushT environment and dataset. - -If you haven't already read our tutorial on the [training script and configuration tooling](../4_train_policy_with_script.md) please do so prior to tackling this tutorial. - -Let's get started! - -Suppose we want to train ACT for PushT. Well, there are aspects of the ACT configuration that are specific to the ALOHA environments, and these happen to be incompatible with PushT. Therefore, trying to run the following will almost certainly raise an exception of sorts (eg: feature dimension mismatch): - -```bash -python lerobot/scripts/train.py policy=act env=pusht dataset_repo_id=lerobot/pusht -``` - -We need to adapt the parameters of the ACT policy configuration to the PushT environment. The most important ones are the image keys. - -ALOHA's datasets and environments typically use a variable number of cameras. In `lerobot/configs/policy/act.yaml` you may notice two relevant sections. Here we show you the minimal diff needed to adjust to PushT: - -```diff -override_dataset_stats: -- observation.images.top: -+ observation.image: - # stats from imagenet, since we use a pretrained vision model - mean: [[[0.485]], [[0.456]], [[0.406]]] # (c,1,1) - std: [[[0.229]], [[0.224]], [[0.225]]] # (c,1,1) - -policy: - input_shapes: -- observation.images.top: [3, 480, 640] -+ observation.image: [3, 96, 96] - observation.state: ["${env.state_dim}"] - output_shapes: - action: ["${env.action_dim}"] - - input_normalization_modes: -- observation.images.top: mean_std -+ observation.image: mean_std - observation.state: min_max - output_normalization_modes: - action: min_max -``` - -Here we've accounted for the following: -- PushT uses "observation.image" for its image key. -- PushT provides smaller images. - -_Side note: technically we could override these via the CLI, but with many changes it gets a bit messy, and we also have a bit of a challenge in that we're using `.` in our observation keys which is treated by Hydra as a hierarchical separator_. - -For your convenience, we provide [`act_pusht.yaml`](./act_pusht.yaml) in this directory. It contains the diff above, plus some other (optional) ones that are explained within. Please copy it into `lerobot/configs/policy` (remember from a [previous tutorial](../4_train_policy_with_script.md) that Hydra will look in the `lerobot/configs` directory). Now try running the following. - - -```bash -python lerobot/scripts/train.py policy=act_pusht env=pusht -``` - -Notice that this is much the same as the command that failed at the start of the tutorial, only: -- Now we are using `policy=act_pusht` to point to our new configuration file. -- We can drop `dataset_repo_id=lerobot/pusht` as the change is incorporated in our new configuration file. - -Hurrah! You're now training ACT for the PushT environment. - ---- - -The bottom line of this tutorial is that when training policies for different environments and datasets you will need to understand what parts of the policy configuration are specific to those and make changes accordingly. - -Happy coding! 🤗 diff --git a/tests/test_examples.py b/tests/test_examples.py index 8bb3f8eb..0a6ce422 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -45,11 +45,11 @@ def test_example_1(): @require_package("gym_pusht") -def test_examples_2_3_calculate_validation_loss(): +def test_examples_basic2_basic3_advanced1(): """ Train a model with example 3, check the outputs. Evaluate the trained model with example 2, check the outputs. - Calculate the validation loss with example 4, check the outputs. + Calculate the validation loss with advanced example 1, check the outputs. """ ### Test example 3 @@ -97,7 +97,7 @@ def test_examples_2_3_calculate_validation_loss(): assert Path("outputs/eval/example_pusht_diffusion/rollout.mp4").exists() ## Test example 4 - file_contents = _read_file("examples/advanced/calculate_validation_loss.py") + file_contents = _read_file("examples/advanced/2_calculate_validation_loss.py") # Run on a single example from the last episode, use CPU, and use the local model. file_contents = _find_and_replace(