hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

fix(DiffusionPolicy): Fix bug where training without image features would crash with exception, fix environment state docs (#1617)

Browse Source 
* Fix bug in diffusion config validation when not using image features

* Fix DiffusionPolicy docstring about shape of env state
This commit is contained in:
Abhay Deshpande
2025-07-29 04:40:16 -07:00
committed by GitHub
parent c14ab9e97b
commit 5695432142
2 changed files with 9 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/lerobot/policies/diffusion/configuration_diffusion.py
View File

@@ -217,12 +217,13 @@ class DiffusionConfig(PreTrainedConfig):
)
# Check that all input images have the same shape.
first_image_key, first_image_ft = next(iter(self.image_features.items()))
for key, image_ft in self.image_features.items():
if image_ft.shape != first_image_ft.shape:
raise ValueError(
f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
)
if len(self.image_features) > 0:
first_image_key, first_image_ft = next(iter(self.image_features.items()))
for key, image_ft in self.image_features.items():
if image_ft.shape != first_image_ft.shape:
raise ValueError(
f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
)
@property
def observation_delta_indices(self) -> list:

4
src/lerobot/policies/diffusion/modeling_diffusion.py
View File

@@ -288,7 +288,7 @@ class DiffusionModel(nn.Module):
"observation.images": (B, n_obs_steps, num_cameras, C, H, W)
AND/OR
"observation.environment_state": (B, environment_dim)
"observation.environment_state": (B, n_obs_steps, environment_dim)
}
"""
batch_size, n_obs_steps = batch["observation.state"].shape[:2]
@@ -315,7 +315,7 @@ class DiffusionModel(nn.Module):
"observation.images": (B, n_obs_steps, num_cameras, C, H, W)
AND/OR
"observation.environment_state": (B, environment_dim)
"observation.environment_state": (B, n_obs_steps, environment_dim)
"action": (B, horizon, action_dim)
"action_is_pad": (B, horizon)