Fix n_obs_steps handling to ensure correct channel shapes

[shay-ff]

🐛 Fix: Training crash when using --policy.n_obs_steps=1

Summary

Fixes a bug where training crashed when --policy.n_obs_steps=1.
Previously, the dataset returned tensors with the wrong number of channels, causing a RuntimeError in the encoder.

Changes

• Converts single-frame images from [H, W, C] → [C, H, W]
• Converts multi-frame images from [N, H, W, C] → [N*C, H, W]
• Applies image_transforms per frame
• Works for any positive n_obs_steps

🧪 How it was tested

• Verified that training starts successfully with both:
  • n_obs_steps=1
  • n_obs_steps>1
• Checked that output image tensors have correct dimensions:
  [n_obs_steps*3, H, W]
• Confirmed that no other functionality is affected.

🔍 How to checkout & try

# Install dependencies if needed
pip install -r requirements.txt

# Train with n_obs_steps=1 (should not crash)
lerobot-train --policy.n_obs_steps=1

# Train with n_obs_steps=2 (should also work)
lerobot-train --policy.n_obs_steps=2

Fixes #2277

[Copilot]

Pull Request Overview

Fixes a bug where training crashed when using n_obs_steps=1 due to incorrect tensor shape handling for image observations. The code now properly converts image tensors from HWC to CHW format and handles both single-frame and multi-frame cases.

• Adds tensor shape conversion logic to handle both [H, W, C] and [N, H, W, C] formats
• Applies image transforms per-frame for multi-frame observations
• Reshapes multi-frame tensors to [N*C, H, W] format expected by the encoder

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[shay-ff]

Hi @jadechoghari this PR addresses issue #2277 and applies all suggested fixes.
Could you please review when you have a moment? Thanks!