lerobot/common/policies/pi0/modeling_pi0.py [71:218]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - def create_sinusoidal_pos_embedding( time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu" ) -> Tensor: """Computes sine-cosine positional embedding vectors for scalar positions.""" if dimension % 2 != 0: raise ValueError(f"dimension ({dimension}) must be divisible by 2") if time.ndim != 1: raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") dtype = get_safe_dtype(torch.float64, device.type) fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) period = min_period * (max_period / min_period) ** fraction # Compute the outer product scaling_factor = 1.0 / period * 2 * math.pi sin_input = scaling_factor[None, :] * time[:, None] pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) return pos_emb def sample_beta(alpha, beta, bsize, device): gamma1 = torch.empty((bsize,), device=device).uniform_(0, 1).pow(1 / alpha) gamma2 = torch.empty((bsize,), device=device).uniform_(0, 1).pow(1 / beta) return gamma1 / (gamma1 + gamma2) def make_att_2d_masks(pad_masks, att_masks): """Copied from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to setup several types of attention, for example: [[1 1 1 1 1 1]]: pure causal attention. [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour. [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a block can attend all previous blocks and all tokens on the same block. Args: input_mask: bool[B, N] true if its part of the input, false if padding. mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on it and 0 where it shares the same attention mask as the previous token. """ if att_masks.ndim != 2: raise ValueError(att_masks.ndim) if pad_masks.ndim != 2: raise ValueError(pad_masks.ndim) cumsum = torch.cumsum(att_masks, dim=1) att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] att_2d_masks = att_2d_masks & pad_2d_masks return att_2d_masks def resize_with_pad(img, width, height, pad_value=-1): # assume no-op when width height fits already if img.ndim != 4: raise ValueError(f"(b,c,h,w) expected, but {img.shape}") cur_height, cur_width = img.shape[2:] ratio = max(cur_width / width, cur_height / height) resized_height = int(cur_height / ratio) resized_width = int(cur_width / ratio) resized_img = F.interpolate( img, size=(resized_height, resized_width), mode="bilinear", align_corners=False ) pad_height = max(0, int(height - resized_height)) pad_width = max(0, int(width - resized_width)) # pad on left and top of image padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value) return padded_img def pad_vector(vector, new_dim): """Can be (batch_size x sequence_length x features_dimension) or (batch_size x features_dimension) """ if vector.shape[-1] == new_dim: return vector shape = list(vector.shape) current_dim = shape[-1] shape[-1] = new_dim new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device) new_vector[..., :current_dim] = vector return new_vector def normalize(x, min_val, max_val): return (x - min_val) / (max_val - min_val) def unnormalize(x, min_val, max_val): return x * (max_val - min_val) + min_val def safe_arcsin(value): # This ensures that the input stays within # [−1,1] to avoid invalid values for arcsin return torch.arcsin(torch.clamp(value, -1.0, 1.0)) def aloha_gripper_to_angular(value): # Aloha transforms the gripper positions into a linear space. The following code # reverses this transformation to be consistent with pi0 which is pretrained in # angular space. # # These values are coming from the Aloha code: # PUPPET_GRIPPER_POSITION_OPEN, PUPPET_GRIPPER_POSITION_CLOSED value = unnormalize(value, min_val=0.01844, max_val=0.05800) # This is the inverse of the angular to linear transformation inside the Interbotix code. def linear_to_radian(linear_position, arm_length, horn_radius): value = (horn_radius**2 + linear_position**2 - arm_length**2) / (2 * horn_radius * linear_position) return safe_arcsin(value) # The constants are taken from the Interbotix code. value = linear_to_radian(value, arm_length=0.036, horn_radius=0.022) # Normalize to [0, 1]. # The values 0.4 and 1.5 were measured on an actual Trossen robot. return normalize(value, min_val=0.4, max_val=1.5) def aloha_gripper_from_angular(value): # Convert from the gripper position used by pi0 to the gripper position that is used by Aloha. # Note that the units are still angular but the range is different. # The values 0.4 and 1.5 were measured on an actual Trossen robot. value = unnormalize(value, min_val=0.4, max_val=1.5) # These values are coming from the Aloha code: # PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE return normalize(value, min_val=-0.6213, max_val=1.4910) def aloha_gripper_from_angular_inv(value): # Directly inverts the gripper_from_angular function. value = unnormalize(value, min_val=-0.6213, max_val=1.4910) return normalize(value, min_val=0.4, max_val=1.5) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - lerobot/common/policies/smolvla/modeling_smolvla.py [176:323]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - def create_sinusoidal_pos_embedding( time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu" ) -> Tensor: """Computes sine-cosine positional embedding vectors for scalar positions.""" if dimension % 2 != 0: raise ValueError(f"dimension ({dimension}) must be divisible by 2") if time.ndim != 1: raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") dtype = get_safe_dtype(torch.float64, device.type) fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) period = min_period * (max_period / min_period) ** fraction # Compute the outer product scaling_factor = 1.0 / period * 2 * math.pi sin_input = scaling_factor[None, :] * time[:, None] pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) return pos_emb def sample_beta(alpha, beta, bsize, device): gamma1 = torch.empty((bsize,), device=device).uniform_(0, 1).pow(1 / alpha) gamma2 = torch.empty((bsize,), device=device).uniform_(0, 1).pow(1 / beta) return gamma1 / (gamma1 + gamma2) def make_att_2d_masks(pad_masks, att_masks): """Copied from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to setup several types of attention, for example: [[1 1 1 1 1 1]]: pure causal attention. [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour. [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a block can attend all previous blocks and all tokens on the same block. Args: input_mask: bool[B, N] true if its part of the input, false if padding. mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on it and 0 where it shares the same attention mask as the previous token. """ if att_masks.ndim != 2: raise ValueError(att_masks.ndim) if pad_masks.ndim != 2: raise ValueError(pad_masks.ndim) cumsum = torch.cumsum(att_masks, dim=1) att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] att_2d_masks = att_2d_masks & pad_2d_masks return att_2d_masks def resize_with_pad(img, width, height, pad_value=-1): # assume no-op when width height fits already if img.ndim != 4: raise ValueError(f"(b,c,h,w) expected, but {img.shape}") cur_height, cur_width = img.shape[2:] ratio = max(cur_width / width, cur_height / height) resized_height = int(cur_height / ratio) resized_width = int(cur_width / ratio) resized_img = F.interpolate( img, size=(resized_height, resized_width), mode="bilinear", align_corners=False ) pad_height = max(0, int(height - resized_height)) pad_width = max(0, int(width - resized_width)) # pad on left and top of image padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value) return padded_img def pad_vector(vector, new_dim): """Can be (batch_size x sequence_length x features_dimension) or (batch_size x features_dimension) """ if vector.shape[-1] == new_dim: return vector shape = list(vector.shape) current_dim = shape[-1] shape[-1] = new_dim new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device) new_vector[..., :current_dim] = vector return new_vector def normalize(x, min_val, max_val): return (x - min_val) / (max_val - min_val) def unnormalize(x, min_val, max_val): return x * (max_val - min_val) + min_val def safe_arcsin(value): # This ensures that the input stays within # [−1,1] to avoid invalid values for arcsin return torch.arcsin(torch.clamp(value, -1.0, 1.0)) def aloha_gripper_to_angular(value): # Aloha transforms the gripper positions into a linear space. The following code # reverses this transformation to be consistent with smolvla which is pretrained in # angular space. # # These values are coming from the Aloha code: # PUPPET_GRIPPER_POSITION_OPEN, PUPPET_GRIPPER_POSITION_CLOSED value = unnormalize(value, min_val=0.01844, max_val=0.05800) # This is the inverse of the angular to linear transformation inside the Interbotix code. def linear_to_radian(linear_position, arm_length, horn_radius): value = (horn_radius**2 + linear_position**2 - arm_length**2) / (2 * horn_radius * linear_position) return safe_arcsin(value) # The constants are taken from the Interbotix code. value = linear_to_radian(value, arm_length=0.036, horn_radius=0.022) # Normalize to [0, 1]. # The values 0.4 and 1.5 were measured on an actual Trossen robot. return normalize(value, min_val=0.4, max_val=1.5) def aloha_gripper_from_angular(value): # Convert from the gripper position used by smolvla to the gripper position that is used by Aloha. # Note that the units are still angular but the range is different. # The values 0.4 and 1.5 were measured on an actual Trossen robot. value = unnormalize(value, min_val=0.4, max_val=1.5) # These values are coming from the Aloha code: # PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE return normalize(value, min_val=-0.6213, max_val=1.4910) def aloha_gripper_from_angular_inv(value): # Directly inverts the gripper_from_angular function. value = unnormalize(value, min_val=-0.6213, max_val=1.4910) return normalize(value, min_val=0.4, max_val=1.5) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -