From 5efea7a73be048a78898dd946362fd3c3a62bff4 Mon Sep 17 00:00:00 2001 From: me Date: Wed, 25 Mar 2026 21:40:35 -0700 Subject: [PATCH 1/2] Fixing actual controller action alignment within padded action tensor. Dimensions conversion 21->25->21. --- nitrogen/mm_tokenizers.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nitrogen/mm_tokenizers.py b/nitrogen/mm_tokenizers.py index e1e6504..2f8a610 100644 --- a/nitrogen/mm_tokenizers.py +++ b/nitrogen/mm_tokenizers.py @@ -157,6 +157,8 @@ def _prepare_action(self, data: dict): ), f"Action dim {n_action_dims} exceeds max allowed {self.max_action_dim}." # Pad the channel dimension + # [batch, 18, 21] -> [batch, 18, 25] badded by zeros on the right. + # Action data is left aligned here [btn[0], btn[1], ..., btn[16], j_left[0], j_left[1], j_right[0], j_right[1], 0, 0, 0, 0] actions = np.pad(actions, ((0, 0), (0, self.max_action_dim - n_action_dims)), "constant") # Create mask: [T, max_action_dim] @@ -233,6 +235,10 @@ def pack_actions(self, buttons, j_left, j_right): return action def unpack_actions(self, actions): + # Action data is still left aligned here [btn[0], btn[1], ..., btn[16], j_left[0], j_left[1], j_right[0], j_right[1], 0, 0, 0, 0]. + # And the dimensions are [batch, 18, 25], not [batch, 18, 21]. + # Need to remove padding before using negative indices for extracting j_left, j_right. + actions = actions[:, :, :21] if self.old_layout: # Unpack the actions into j_left, j_right, buttons j_left = actions[:, :, :2] From 62dd42ae2a78c9a80590aba0eb0ee8db07387f1e Mon Sep 17 00:00:00 2001 From: me Date: Wed, 25 Mar 2026 21:54:31 -0700 Subject: [PATCH 2/2] The issue with alignment is when self.old_layout is False. --- nitrogen/mm_tokenizers.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/nitrogen/mm_tokenizers.py b/nitrogen/mm_tokenizers.py index 2f8a610..f393216 100644 --- a/nitrogen/mm_tokenizers.py +++ b/nitrogen/mm_tokenizers.py @@ -235,16 +235,16 @@ def pack_actions(self, buttons, j_left, j_right): return action def unpack_actions(self, actions): - # Action data is still left aligned here [btn[0], btn[1], ..., btn[16], j_left[0], j_left[1], j_right[0], j_right[1], 0, 0, 0, 0]. - # And the dimensions are [batch, 18, 25], not [batch, 18, 21]. - # Need to remove padding before using negative indices for extracting j_left, j_right. - actions = actions[:, :, :21] if self.old_layout: # Unpack the actions into j_left, j_right, buttons j_left = actions[:, :, :2] j_right = actions[:, :, 2:4] buttons = actions[:, :, 4:] else: + # Action data is still left aligned here [btn[0], btn[1], ..., btn[16], j_left[0], j_left[1], j_right[0], j_right[1], 0, 0, 0, 0]. + # And the dimensions are [batch, 18, 25], not [batch, 18, 21]. + # Need to remove padding before using negative indices for extracting j_left, j_right. + actions = actions[:, :, :21] # Unpack the actions into j_left, j_right, buttons buttons = actions[:, :, :-4] j_left = actions[:, :, -4:-2]