From 5361346bece1f20535bf5d331793a1d5ba1d79cc Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 11:25:26 +0200 Subject: [PATCH 1/6] Do not add model prefix to normalization --- src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py | 5 ++++- src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py index 9ff71152a..1f507c75d 100644 --- a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py @@ -919,7 +919,10 @@ class PI05OpenPIPolicy(PreTrainedPolicy): remap_count = 0 for key, value in fixed_state_dict.items(): - if not key.startswith("model."): + if not key.startswith("model.") and not any( + key.startswith(prefix) + for prefix in ["normalize_inputs.", "normalize_targets.", "unnormalize_outputs."] + ): new_key = f"model.{key}" remapped_state_dict[new_key] = value remap_count += 1 diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 549dc0a9b..1fdb6048b 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -938,7 +938,10 @@ class PI0OpenPIPolicy(PreTrainedPolicy): remap_count = 0 for key, value in fixed_state_dict.items(): - if not key.startswith("model."): + if not key.startswith("model.") and not any( + key.startswith(prefix) + for prefix in ["normalize_inputs.", "normalize_targets.", "unnormalize_outputs."] + ): new_key = f"model.{key}" remapped_state_dict[new_key] = value remap_count += 1 From b9df1a4ac5d12335ea25748ba4f1c10a88a8c20d Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 13:08:41 +0200 Subject: [PATCH 2/6] use same name for action and state dim as lerobot pi0 and remove fixed image keys --- .../pi05_openpi/configuration_pi05openpi.py | 39 ++++++------------- .../pi05_openpi/modeling_pi05openpi.py | 16 ++++---- .../pi0_openpi/configuration_pi0openpi.py | 39 ++++++------------- .../policies/pi0_openpi/modeling_pi0openpi.py | 16 ++++---- 4 files changed, 38 insertions(+), 72 deletions(-) diff --git a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py index 4566fe0e7..acf2b7b6e 100644 --- a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py @@ -16,7 +16,7 @@ from dataclasses import dataclass, field from lerobot.configs.policies import PreTrainedConfig -from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature +from lerobot.configs.types import NormalizationMode from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig @@ -36,23 +36,20 @@ class PI05OpenPIConfig(PreTrainedConfig): n_obs_steps: int = 1 chunk_size: int = 50 # Number of action steps to predict, in openpi called "action_horizon" n_action_steps: int = 50 # Number of action steps to execute - action_dim: int = 32 # Action dimension (will be padded to 32) - state_dim: int = 32 # State dimension (will be padded to 32) + + # Shorter state and action vectors will be padded to these dimensions + max_state_dim: int = 32 # State dimension (will be padded to 32) + max_action_dim: int = 32 # Action dimension (will be padded to 32) # Flow matching parameters: see openpi `PI0Pytorch` num_inference_steps: int = 10 # Number of denoising steps during inference time_sampling_beta_alpha: float = 1.5 # Beta distribution alpha parameter for time sampling time_sampling_beta_beta: float = 1.0 # Beta distribution beta parameter for time sampling min_period: float = 4e-3 # Min period for sinusoidal positional encoding - max_period: float = 4.0 # Max period for sinusoidal positional encodingis my + max_period: float = 4.0 # Max period for sinusoidal positional encoding # Image preprocessing image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py` - image_keys: tuple[str, ...] = ( - "observation.images.base_0_rgb", - "observation.images.left_wrist_0_rgb", - "observation.images.right_wrist_0_rgb", - ) # Normalization normalization_mapping: dict[str, NormalizationMode] = field( @@ -103,26 +100,12 @@ class PI05OpenPIConfig(PreTrainedConfig): def validate_features(self) -> None: """Validate and set up input/output features.""" - # Add image features - for key in self.image_keys: - if key not in self.input_features: - self.input_features[key] = PolicyFeature( - type=FeatureType.VISUAL, - shape=(3, 224, 224), # Default shape, will be resized - ) + # Image features are now handled dynamically through dataset configuration + # No need to auto-add hardcoded image keys - # Ensure state and action features exist - if "observation.state" not in self.input_features: - self.input_features["observation.state"] = PolicyFeature( - type=FeatureType.STATE, - shape=(self.state_dim,), - ) - - if "action" not in self.output_features: - self.output_features["action"] = PolicyFeature( - type=FeatureType.ACTION, - shape=(self.action_dim,), - ) + # State and action features are also handled dynamically through dataset configuration + # The actual dimensions come from the feature shapes, max dimensions are used for padding only + pass def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig( diff --git a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py index 1f507c75d..455640213 100644 --- a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py @@ -503,8 +503,8 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` precision=config.dtype, ) - self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) - self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim) + self.action_in_proj = nn.Linear(config.max_action_dim, action_expert_config.width) + self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim) self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width) self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) @@ -739,8 +739,8 @@ $(python -c "import transformers, os; print(os.path.dirname(transformers.__file_ actions_shape = ( bsize, self.config.chunk_size, - self.config.action_dim, - ) # Use config action_dim for internal processing + self.config.max_action_dim, + ) # Use config max_action_dim for internal processing noise = self.sample_noise(actions_shape, device) prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( @@ -1235,12 +1235,12 @@ class PI05OpenPIPolicy(PreTrainedPolicy): def prepare_state(self, batch): # see lerobot pi0 `prepare_state` (exact copy) """Pad state""" - state = pad_vector(batch[OBS_STATE], self.config.state_dim) + state = pad_vector(batch[OBS_STATE], self.config.max_state_dim) return state def prepare_action(self, batch): # see lerobot pi0 `prepare_action` (exact copy) """Pad action""" - actions = pad_vector(batch[ACTION], self.config.action_dim) + actions = pad_vector(batch[ACTION], self.config.max_action_dim) return actions @torch.no_grad() @@ -1294,8 +1294,8 @@ class PI05OpenPIPolicy(PreTrainedPolicy): losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions) # Truncate losses to actual action dimensions - if self.config.action_dim < 32: - losses = losses[:, :, : self.config.action_dim] + original_action_dim = self.config.output_features[ACTION].shape[0] + losses = losses[:, :, :original_action_dim] loss = losses.mean() diff --git a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py index 59d0085ae..973638fae 100644 --- a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py @@ -16,7 +16,7 @@ from dataclasses import dataclass, field from lerobot.configs.policies import PreTrainedConfig -from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature +from lerobot.configs.types import NormalizationMode from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig @@ -33,23 +33,20 @@ class PI0OpenPIConfig(PreTrainedConfig): n_obs_steps: int = 1 chunk_size: int = 50 # Number of action steps to predict, in openpi called "action_horizon" n_action_steps: int = 50 # Number of action steps to execute - action_dim: int = 32 # Action dimension (will be padded to 32) - state_dim: int = 32 # State dimension (will be padded to 32) + + # Shorter state and action vectors will be padded to these dimensions + max_state_dim: int = 32 # State dimension (will be padded to 32) + max_action_dim: int = 32 # Action dimension (will be padded to 32) # Flow matching parameters: see openpi `PI0Pytorch` num_inference_steps: int = 10 # Number of denoising steps during inference time_sampling_beta_alpha: float = 1.5 # Beta distribution alpha parameter for time sampling time_sampling_beta_beta: float = 1.0 # Beta distribution beta parameter for time sampling min_period: float = 4e-3 # Min period for sinusoidal positional encoding - max_period: float = 4.0 # Max period for sinusoidal positional encodingis my + max_period: float = 4.0 # Max period for sinusoidal positional encoding # Image preprocessing image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py` - image_keys: tuple[str, ...] = ( - "observation.images.base_0_rgb", - "observation.images.left_wrist_0_rgb", - "observation.images.right_wrist_0_rgb", - ) # Normalization normalization_mapping: dict[str, NormalizationMode] = field( @@ -100,26 +97,12 @@ class PI0OpenPIConfig(PreTrainedConfig): def validate_features(self) -> None: """Validate and set up input/output features.""" - # Add image features - for key in self.image_keys: - if key not in self.input_features: - self.input_features[key] = PolicyFeature( - type=FeatureType.VISUAL, - shape=(3, 224, 224), # Default shape, will be resized - ) + # Image features are now handled dynamically through dataset configuration + # No need to auto-add hardcoded image keys - # Ensure state and action features exist - if "observation.state" not in self.input_features: - self.input_features["observation.state"] = PolicyFeature( - type=FeatureType.STATE, - shape=(self.state_dim,), - ) - - if "action" not in self.output_features: - self.output_features["action"] = PolicyFeature( - type=FeatureType.ACTION, - shape=(self.action_dim,), - ) + # State and action features are also handled dynamically through dataset configuration + # The actual dimensions come from the feature shapes, max dimensions are used for padding only + pass def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig( diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 1fdb6048b..3f82cfbcf 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -503,10 +503,10 @@ class PI0Pytorch(nn.Module): # see openpi `PI0Pytorch` precision=config.dtype, ) - self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) - self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim) + self.action_in_proj = nn.Linear(config.max_action_dim, action_expert_config.width) + self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim) - self.state_proj = nn.Linear(config.state_dim, action_expert_config.width) + self.state_proj = nn.Linear(config.max_state_dim, action_expert_config.width) self.action_time_mlp_in = nn.Linear(2 * action_expert_config.width, action_expert_config.width) self.action_time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) @@ -758,8 +758,8 @@ $(python -c "import transformers, os; print(os.path.dirname(transformers.__file_ actions_shape = ( bsize, self.config.chunk_size, - self.config.action_dim, - ) # Use config action_dim for internal processing + self.config.max_action_dim, + ) # Use config max_action_dim for internal processing noise = self.sample_noise(actions_shape, device) prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( @@ -1250,12 +1250,12 @@ class PI0OpenPIPolicy(PreTrainedPolicy): def prepare_state(self, batch): # see lerobot pi0 `prepare_state` (exact copy) """Pad state""" - state = pad_vector(batch[OBS_STATE], self.config.state_dim) + state = pad_vector(batch[OBS_STATE], self.config.max_state_dim) return state def prepare_action(self, batch): # see lerobot pi0 `prepare_action` (exact copy) """Pad action""" - actions = pad_vector(batch[ACTION], self.config.action_dim) + actions = pad_vector(batch[ACTION], self.config.max_action_dim) return actions @torch.no_grad() @@ -1314,7 +1314,7 @@ class PI0OpenPIPolicy(PreTrainedPolicy): loss = losses.mean() loss_dict = { - "loss": loss.item(), + "l2_loss": loss.item(), "loss_per_dim": losses.mean(dim=[0, 1]).detach().cpu().numpy().tolist(), } From af0676f99e79d15472c542b0f41c9e7cc9497049 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 14:27:07 +0200 Subject: [PATCH 3/6] load from pretrained_path --- src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py | 3 +++ src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py | 4 ++-- src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py | 3 +++ src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py | 4 ++-- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py index acf2b7b6e..2902f8216 100644 --- a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py @@ -66,6 +66,9 @@ class PI05OpenPIConfig(PreTrainedConfig): compile_mode: str = "max-autotune" # Torch compile mode device: str | None = None # Device to use for the model (None = auto-detect) + # Pretrained model loading + pretrained_path: str | None = None # Path or repo_id to load pretrained weights from + # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` optimizer_betas: tuple[float, float] = (0.9, 0.95) diff --git a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py index 455640213..fb2a4ea49 100644 --- a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py @@ -875,8 +875,8 @@ class PI05OpenPIPolicy(PreTrainedPolicy): if pretrained_name_or_path is None: raise ValueError("pretrained_name_or_path is required") - # Create default config - config = cls.config_class() + # Use provided config if available, otherwise create default config + config = kwargs.get("config", cls.config_class()) # Initialize model without loading weights # Check if dataset_stats were provided in kwargs diff --git a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py index 973638fae..52e587abe 100644 --- a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py @@ -63,6 +63,9 @@ class PI0OpenPIConfig(PreTrainedConfig): compile_mode: str = "max-autotune" # Torch compile mode device: str | None = None # Device to use for the model (None = auto-detect) + # Pretrained model loading + pretrained_path: str | None = None # Path or repo_id to load pretrained weights from + # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` optimizer_betas: tuple[float, float] = (0.9, 0.95) diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 3f82cfbcf..4bb1c83d9 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -894,8 +894,8 @@ class PI0OpenPIPolicy(PreTrainedPolicy): if pretrained_name_or_path is None: raise ValueError("pretrained_name_or_path is required") - # Create default config - config = cls.config_class() + # Use provided config if available, otherwise create default config + config = kwargs.get("config", cls.config_class()) # Initialize model without loading weights # Check if dataset_stats were provided in kwargs From 6db39cad58cc745929b91fd93c2f08ad951104b3 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 14:43:09 +0200 Subject: [PATCH 4/6] temp: hardcode base model --- src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py | 2 +- src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py index 2902f8216..03938e570 100644 --- a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py @@ -67,7 +67,7 @@ class PI05OpenPIConfig(PreTrainedConfig): device: str | None = None # Device to use for the model (None = auto-detect) # Pretrained model loading - pretrained_path: str | None = None # Path or repo_id to load pretrained weights from + pretrained_path: str | None = "pepijn223/pi0_base_fp32" # Path or repo_id to load pretrained weights from # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` diff --git a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py index 52e587abe..1f8583849 100644 --- a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py @@ -64,7 +64,7 @@ class PI0OpenPIConfig(PreTrainedConfig): device: str | None = None # Device to use for the model (None = auto-detect) # Pretrained model loading - pretrained_path: str | None = None # Path or repo_id to load pretrained weights from + pretrained_path: str | None = "pepijn223/pi0_base_fp32" # Path or repo_id to load pretrained weights from # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` From 7d9b469eeea30bfef9083a3bded09ff61c6f211e Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 14:50:43 +0200 Subject: [PATCH 5/6] fix override self.pretrained_path = None overwrite --- src/lerobot/configs/policies.py | 4 +++- src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py | 3 --- src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py | 3 --- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/src/lerobot/configs/policies.py b/src/lerobot/configs/policies.py index f5fa727cf..418daa0e7 100644 --- a/src/lerobot/configs/policies.py +++ b/src/lerobot/configs/policies.py @@ -72,9 +72,11 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): tags: list[str] | None = None # Add tags to your policy on the hub. license: str | None = None + # Either the repo ID of a model hosted on the Hub or a path to a directory containing weights + # saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch. + pretrained_path: str | None = None def __post_init__(self): - self.pretrained_path = None if not self.device or not is_torch_device_available(self.device): auto_device = auto_select_torch_device() logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.") diff --git a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py index 03938e570..acf2b7b6e 100644 --- a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py @@ -66,9 +66,6 @@ class PI05OpenPIConfig(PreTrainedConfig): compile_mode: str = "max-autotune" # Torch compile mode device: str | None = None # Device to use for the model (None = auto-detect) - # Pretrained model loading - pretrained_path: str | None = "pepijn223/pi0_base_fp32" # Path or repo_id to load pretrained weights from - # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` optimizer_betas: tuple[float, float] = (0.9, 0.95) diff --git a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py index 1f8583849..973638fae 100644 --- a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py @@ -63,9 +63,6 @@ class PI0OpenPIConfig(PreTrainedConfig): compile_mode: str = "max-autotune" # Torch compile mode device: str | None = None # Device to use for the model (None = auto-detect) - # Pretrained model loading - pretrained_path: str | None = "pepijn223/pi0_base_fp32" # Path or repo_id to load pretrained weights from - # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` optimizer_betas: tuple[float, float] = (0.9, 0.95) From d0d714be4784870bee97c8226f089670a0323b21 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 16:15:29 +0200 Subject: [PATCH 6/6] rename to loss --- src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 4bb1c83d9..8cfa1743b 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -1314,7 +1314,7 @@ class PI0OpenPIPolicy(PreTrainedPolicy): loss = losses.mean() loss_dict = { - "l2_loss": loss.item(), + "loss": loss.item(), "loss_per_dim": losses.mean(dim=[0, 1]).detach().cpu().numpy().tolist(), }