more styling fixes:

This commit is contained in:
Jade Choghari
2025-11-18 16:01:41 +01:00
parent 769eb27c87
commit a068618faf
@@ -290,29 +290,25 @@ class LiberoProcessorStep(ObservationProcessorStep):
""" """
if not isinstance(quat, torch.Tensor): if not isinstance(quat, torch.Tensor):
raise TypeError( raise TypeError(f"_quat2axisangle expected a torch.Tensor, got {type(quat)}")
f"_quat2axisangle expected a torch.Tensor, got {type(quat)}"
)
if quat.ndim != 2 or quat.shape[1] != 4: if quat.ndim != 2 or quat.shape[1] != 4:
raise ValueError( raise ValueError(f"_quat2axisangle expected shape (B, 4), got {tuple(quat.shape)}")
f"_quat2axisangle expected shape (B, 4), got {tuple(quat.shape)}"
)
quat = quat.to(dtype=torch.float32) quat = quat.to(dtype=torch.float32)
device = quat.device device = quat.device
B = quat.shape[0] batch_size = quat.shape[0]
w = quat[:, 3].clamp(-1.0, 1.0) w = quat[:, 3].clamp(-1.0, 1.0)
den = torch.sqrt(torch.clamp(1.0 - w * w, min=0.0)) den = torch.sqrt(torch.clamp(1.0 - w * w, min=0.0))
result = torch.zeros((B, 3), device=device) result = torch.zeros((batch_size, 3), device=device)
mask = den > 1e-10 mask = den > 1e-10
if mask.any(): if mask.any():
angle = 2.0 * torch.acos(w[mask]) # (M,) angle = 2.0 * torch.acos(w[mask]) # (M,)
axis = quat[mask, :3] / den[mask].unsqueeze(1) axis = quat[mask, :3] / den[mask].unsqueeze(1)
result[mask] = axis * angle.unsqueeze(1) result[mask] = axis * angle.unsqueeze(1)