fix(train): handle NHWC images in wandb logging

[wadeKeith]

What this PR does

Fixes wandb image logging crash when the dataloader returns images in NHWC format.

Problem

The unified dataloader in PyTorch mode may return observation images in NHWC [B, H, W, C] layout (consistent with JAX/LeRobot conventions). The current wandb logging code unconditionally calls permute(1, 2, 0) assuming NCHW input. When images are already NHWC:

frame shape: [224, 224, 3]          # NHWC
after permute(1, 2, 0): [224, 3, 224]  # broken
concatenate 3 views on dim=1: [224, 9, 224]  # wandb rejects this

ValueError: Un-supported shape for image conversion [224, 9, 224]

Fix

Detect the image layout before permuting:

• If frame.shape[0] is a small channel count (1 or 3) → treat as NCHW, permute to NHWC
• Otherwise → already NHWC, use as-is

Also:

• Concatenate along the width axis (dim=1 in NHWC) for correct horizontal tiling
• Normalise pixel values from [-1, 1] to [0, 1] when needed, avoiding wandb range warnings

Changes

File	Change
scripts/train_pytorch.py	Robust NCHW/NHWC detection + value normalisation in wandb image logging

Checklist

• Backward compatible (NCHW images still handled correctly)
• Minimal change — 1 file
• No new dependencies

Closes #877

[wadeKeith]

Friendly ping — just checking in on this PR to see if you have any feedback. Happy to address any questions! 🙏

[wadeKeith]

@kvablack gentle ping — all CI is green, ready for review. Thanks!

[wadeKeith]

Friendly ping for review. Thanks!