lerobot_train script fails training with lerobot/libero dataset

[ozgraslan]

Ticket Type

🐛 Bug Report (Something isn't working)

Environment & System Info

- LeRobot version: 0.4.4
- Platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39
- Python version: 3.10.18
- Huggingface Hub version: 0.35.3
- Datasets version: 4.1.1
- Numpy version: 2.2.6
- FFmpeg version: N/A
- PyTorch version: 2.7.1+cu126
- Is PyTorch built with CUDA support?: True
- Cuda version: 12.6
- GPU model: NVIDIA GeForce RTX 4060 Laptop GPU
- Using GPU in script?: <fill in>
- lerobot scripts: ['lerobot-calibrate', 'lerobot-dataset-viz', 'lerobot-edit-dataset', 'lerobot-eval', 'lerobot-find-cameras', 'lerobot-find-joint-limits', 'lerobot-find-port', 'lerobot-imgtransform-viz', 'lerobot-info', 'lerobot-record', 'lerobot-replay', 'lerobot-setup-can', 'lerobot-setup-motors', 'lerobot-teleoperate', 'lerobot-train', 'lerobot-train-tokenizer']

Description

Training a policy using the lerobot_train.py script with lerobot/libero dataset gives the following error:

Traceback (most recent call last):
  File "/home/ozgraslan/git/lerobot_test/.lerobot/bin/lerobot-train", line 10, in <module>
    sys.exit(main())
  File "/home/ozgraslan/git/lerobot_test/src/lerobot/scripts/lerobot_train.py", line 532, in main
    train()
  File "/home/ozgraslan/git/lerobot_test/src/lerobot/configs/parser.py", line 233, in wrapper_inner
    response = fn(cfg, *args, **kwargs)
  File "/home/ozgraslan/git/lerobot_test/src/lerobot/scripts/lerobot_train.py", line 216, in train
    dataset = make_dataset(cfg)
  File "/home/ozgraslan/git/lerobot_test/src/lerobot/datasets/factory.py", line 131, in make_dataset
    dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
KeyError: 'observation.images.image'

Context & Reproduction

lerobot-train
--output_dir=/libero_test
--policy.type=diffusion
--policy.repo_id=[username]/diffusion_policy_libero_test
--dataset.repo_id=lerobot/libero
--env.type=libero
--seed=100000
--batch_size=2
--num_workers=-1
--steps=50
--eval_freq=10
--wandb.enable=false

Relevant logs or stack trace

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

No response

[sean1295]

lerobot/libero meta stats are flawed.

Simply instantiate missing stats before assigning them like follows:

    if cfg.dataset.use_imagenet_stats:
        for key in dataset.meta.camera_keys:
            if key not in dataset.meta.stats:
                dataset.meta.stats[key] = {}
            for stats_type, stats in IMAGENET_STATS.items():
                dataset.meta.stats[key][stats_type] = torch.tensor(
                    stats, dtype=torch.float32
                )