# from abc import ABCMeta, abstractmethod, classmethod from abc import abstractmethod from collections import OrderedDict import json import os import xml.etree.ElementTree as ET import numpy as np import scipy.ndimage.filters as filters import torch from gear_sonic.isaac_utils.rotations import ( quat_identity_like, quat_inverse, quat_mul_norm, transform_from_rotation_translation, transform_mul, transform_rotation, transform_translation, ) class NumpyEncoder(json.JSONEncoder): """Special json encoder for numpy types""" def default(self, obj): if isinstance( obj, np.int_ | np.intc | np.intp | np.int8 | np.int16 | np.int32 | np.int64 | np.uint8 | np.uint16 | np.uint32 | np.uint64, ): return int(obj) elif isinstance(obj, np.float_ | np.float16 | np.float32 | np.float64): return float(obj) elif isinstance(obj, np.ndarray): return {"__ndarray__": obj.tolist(), "dtype": str(obj.dtype), "shape": obj.shape} return json.JSONEncoder.default(self, obj) def json_numpy_obj_hook(dct): if isinstance(dct, dict) and "__ndarray__" in dct: data = np.asarray(dct["__ndarray__"], dtype=dct["dtype"]) return data.reshape(dct["shape"]) return dct class Serializable: """Implementation to read/write to file. All class the is inherited from this class needs to implement to_dict() and from_dict() """ @classmethod def from_dict(cls, dict_repr, *args, **kwargs): """Read the object from an ordered dictionary :param dict_repr: the ordered dictionary that is used to construct the object :type dict_repr: OrderedDict :param args, kwargs: the arguments that need to be passed into from_dict() :type args, kwargs: additional arguments """ pass @abstractmethod def to_dict(self): """Construct an ordered dictionary from the object :rtype: OrderedDict """ pass @classmethod def from_file(cls, path, *args, **kwargs): """Read the object from a file (either .npy or .json) :param path: path of the file :type path: string :param args, kwargs: the arguments that need to be passed into from_dict() :type args, kwargs: additional arguments """ if path.endswith(".json"): with open(path) as f: d = json.load(f, object_hook=json_numpy_obj_hook) elif path.endswith(".npy"): d = np.load(path, allow_pickle=True).item() else: assert False, f"failed to load {cls.__name__} from {path}" assert d["__name__"] == cls.__name__, "the file belongs to {}, not {}".format( d["__name__"], cls.__name__ ) return cls.from_dict(d, *args, **kwargs) def to_file(self, path: str) -> None: """Write the object to a file (either .npy or .json) :param path: path of the file :type path: string """ if os.path.dirname(path) != "" and not os.path.exists(os.path.dirname(path)): os.makedirs(os.path.dirname(path)) d = self.to_dict() d["__name__"] = self.__class__.__name__ if path.endswith(".json"): with open(path, "w") as f: json.dump(d, f, cls=NumpyEncoder, indent=4) elif path.endswith(".npy"): np.save(path, d) class TensorUtils(Serializable): @classmethod def from_dict(cls, dict_repr, *args, **kwargs): """Read the object from an ordered dictionary :param dict_repr: the ordered dictionary that is used to construct the object :type dict_repr: OrderedDict :param kwargs: the arguments that need to be passed into from_dict() :type kwargs: additional arguments """ return torch.from_numpy(dict_repr["arr"].astype(dict_repr["context"]["dtype"])) def to_dict(self): """Construct an ordered dictionary from the object :rtype: OrderedDict """ return NotImplemented def tensor_to_dict(x): """Construct an ordered dictionary from the object :rtype: OrderedDict """ x_np = x.numpy() return {"arr": x_np, "context": {"dtype": x_np.dtype.name}} class SkeletonTree(Serializable): """ A skeleton tree gives a complete description of a rigid skeleton. It describes a tree structure over a list of nodes with their names indicated by strings. Each edge in the tree has a local translation associated with it which describes the distance between the two nodes that it connects. Basic Usage: >>> t = SkeletonTree.from_mjcf(SkeletonTree.__example_mjcf_path__) >>> t SkeletonTree( node_names=['torso', 'front_left_leg', 'aux_1', 'front_left_foot', 'front_right_leg', 'aux_2', 'front_right_foot', 'left_back_leg', 'aux_3', 'left_back_foot', 'right_back_leg', 'aux_4', 'right_back_foot'], parent_indices=tensor([-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11]), local_translation=tensor([[ 0.0000, 0.0000, 0.7500], [ 0.0000, 0.0000, 0.0000], [ 0.2000, 0.2000, 0.0000], [ 0.2000, 0.2000, 0.0000], [ 0.0000, 0.0000, 0.0000], [-0.2000, 0.2000, 0.0000], [-0.2000, 0.2000, 0.0000], [ 0.0000, 0.0000, 0.0000], [-0.2000, -0.2000, 0.0000], [-0.2000, -0.2000, 0.0000], [ 0.0000, 0.0000, 0.0000], [ 0.2000, -0.2000, 0.0000], [ 0.2000, -0.2000, 0.0000]]) ) >>> t.node_names ['torso', 'front_left_leg', 'aux_1', 'front_left_foot', 'front_right_leg', 'aux_2', 'front_right_foot', 'left_back_leg', 'aux_3', 'left_back_foot', 'right_back_leg', 'aux_4', 'right_back_foot'] >>> t.parent_indices tensor([-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11]) >>> t.local_translation tensor([[ 0.0000, 0.0000, 0.7500], [ 0.0000, 0.0000, 0.0000], [ 0.2000, 0.2000, 0.0000], [ 0.2000, 0.2000, 0.0000], [ 0.0000, 0.0000, 0.0000], [-0.2000, 0.2000, 0.0000], [-0.2000, 0.2000, 0.0000], [ 0.0000, 0.0000, 0.0000], [-0.2000, -0.2000, 0.0000], [-0.2000, -0.2000, 0.0000], [ 0.0000, 0.0000, 0.0000], [ 0.2000, -0.2000, 0.0000], [ 0.2000, -0.2000, 0.0000]]) >>> t.parent_of('front_left_leg') 'torso' >>> t.index('front_right_foot') 6 >>> t[2] 'aux_1' """ __example_mjcf_path__ = os.path.join( os.path.dirname(os.path.realpath(__file__)), "tests/ant.xml" ) def __init__(self, node_names, parent_indices, local_translation): """ :param node_names: a list of names for each tree node :type node_names: list[str] :param parent_indices: an int32-typed tensor that represents the edge to its parent.\ -1 represents the root node :type parent_indices: Tensor :param local_translation: a 3d vector that gives local translation information :type local_translation: Tensor """ ln, lp, ll = len(node_names), len(parent_indices), len(local_translation) assert len({ln, lp, ll}) == 1 self._node_names = node_names self._parent_indices = parent_indices.long() self._local_translation = local_translation self._node_indices = {self.node_names[i]: i for i in range(len(self))} def __len__(self): """number of nodes in the skeleton tree""" return len(self.node_names) def __iter__(self): """iterator that iterate through the name of each node""" yield from self.node_names def __getitem__(self, item): """get the name of the node given the index""" return self.node_names[item] def __repr__(self): return ( f"SkeletonTree(\n node_names={self._indent(repr(self.node_names))},\n parent_indices={self._indent(repr(self.parent_indices))}," f"\n local_translation={self._indent(repr(self.local_translation))}\n)" ) def _indent(self, s): return "\n ".join(s.split("\n")) @property def node_names(self): return self._node_names @property def parent_indices(self): return self._parent_indices @property def local_translation(self): return self._local_translation @property def num_joints(self): """number of nodes in the skeleton tree""" return len(self) @classmethod def from_dict(cls, dict_repr, *args, **kwargs): return cls( list(map(str, dict_repr["node_names"])), TensorUtils.from_dict(dict_repr["parent_indices"], *args, **kwargs), TensorUtils.from_dict(dict_repr["local_translation"], *args, **kwargs), ) def to_dict(self): return OrderedDict( [ ("node_names", self.node_names), ("parent_indices", tensor_to_dict(self.parent_indices)), ("local_translation", tensor_to_dict(self.local_translation)), ] ) @classmethod def from_mjcf(cls, path: str) -> "SkeletonTree": """ Parses a mujoco xml scene description file and returns a Skeleton Tree. We use the model attribute at the root as the name of the tree. :param path: :type path: string :return: The skeleton tree constructed from the mjcf file :rtype: SkeletonTree """ tree = ET.parse(path) xml_doc_root = tree.getroot() xml_world_body = xml_doc_root.find("worldbody") if xml_world_body is None: raise ValueError("MJCF parsed incorrectly please verify it.") # assume this is the root xml_body_root = xml_world_body.find("body") if xml_body_root is None: raise ValueError("MJCF parsed incorrectly please verify it.") node_names = [] parent_indices = [] local_translation = [] # recursively adding all nodes into the skel_tree def _add_xml_node(xml_node, parent_index, node_index): node_name = xml_node.attrib.get("name") # parse the local translation into float list pos = np.fromstring(xml_node.attrib.get("pos", "0 0 0"), dtype=float, sep=" ") node_names.append(node_name) parent_indices.append(parent_index) local_translation.append(pos) curr_index = node_index node_index += 1 for next_node in xml_node.findall("body"): node_index = _add_xml_node(next_node, curr_index, node_index) return node_index _add_xml_node(xml_body_root, -1, 0) return cls( node_names, torch.from_numpy(np.array(parent_indices, dtype=np.int32)), torch.from_numpy(np.array(local_translation, dtype=np.float32)), ) def parent_of(self, node_name): """get the name of the parent of the given node :param node_name: the name of the node :type node_name: string :rtype: string """ return self[int(self.parent_indices[self.index(node_name)].item())] def index(self, node_name): """get the index of the node :param node_name: the name of the node :type node_name: string :rtype: int """ return self._node_indices[node_name] def drop_nodes_by_names( self, node_names: list[str], pairwise_translation=None ) -> "SkeletonTree": new_length = len(self) - len(node_names) new_node_names = [] new_local_translation = torch.zeros(new_length, 3, dtype=self.local_translation.dtype) new_parent_indices = torch.zeros(new_length, dtype=self.parent_indices.dtype) parent_indices = self.parent_indices.numpy() new_node_indices: dict = {} new_node_index = 0 for node_index in range(len(self)): if self[node_index] in node_names: continue tb_node_index = parent_indices[node_index] if tb_node_index != -1: local_translation = self.local_translation[node_index, :] while tb_node_index != -1 and self[tb_node_index] in node_names: local_translation += self.local_translation[tb_node_index, :] tb_node_index = parent_indices[tb_node_index] assert tb_node_index != -1, "the root node cannot be dropped" if pairwise_translation is not None: local_translation = pairwise_translation[tb_node_index, node_index, :] else: local_translation = self.local_translation[node_index, :] new_node_names.append(self[node_index]) new_local_translation[new_node_index, :] = local_translation if tb_node_index == -1: new_parent_indices[new_node_index] = -1 else: new_parent_indices[new_node_index] = new_node_indices[self[tb_node_index]] new_node_indices[self[node_index]] = new_node_index new_node_index += 1 return SkeletonTree(new_node_names, new_parent_indices, new_local_translation) def keep_nodes_by_names( self, node_names: list[str], pairwise_translation=None ) -> "SkeletonTree": nodes_to_drop = list(filter(lambda x: x not in node_names, self)) return self.drop_nodes_by_names(nodes_to_drop, pairwise_translation) class SkeletonState(Serializable): """ A skeleton state contains all the information needed to describe a static state of a skeleton. It requires a skeleton tree, local/global rotation at each joint and the root translation. Example: >>> t = SkeletonTree.from_mjcf(SkeletonTree.__example_mjcf_path__) >>> zero_pose = SkeletonState.zero_pose(t) >>> plot_skeleton_state(zero_pose) # can be imported from `.visualization.base_interface` [plot of the ant at zero pose >>> local_rotation = zero_pose.local_rotation.clone() >>> local_rotation[2] = torch.tensor([0, 0, 1, 0]) >>> new_pose = SkeletonState.from_rotation_and_root_translation( ... skeleton_tree=t, ... r=local_rotation, ... t=zero_pose.root_translation, ... is_local=True ... ) >>> new_pose.local_rotation tensor([[0., 0., 0., 1.], [0., 0., 0., 1.], [0., 1., 0., 0.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.]]) >>> plot_skeleton_state(new_pose) # you should be able to see one of ant's leg is bent [plot of the ant with the new pose >>> new_pose.global_rotation # the local rotation is propagated to the global rotation at joint #3 tensor([[0., 0., 0., 1.], [0., 0., 0., 1.], [0., 1., 0., 0.], [0., 1., 0., 0.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.], [0., 0., 0., 1.]]) Global/Local Representation (cont. from the previous example) >>> new_pose.is_local True >>> new_pose.tensor # this will return the local rotation followed by the root translation tensor([0., 0., 0., 1., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0.]) >>> new_pose.tensor.shape # 4 * 13 (joint rotation) + 3 (root translatio torch.Size([55]) >>> new_pose.global_repr().is_local False >>> new_pose.global_repr().tensor # this will return the global rotation followed by the root translation instead tensor([0., 0., 0., 1., 0., 0., 0., 1., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0.]) >>> new_pose.global_repr().tensor.shape # 4 * 13 (joint rotation) + 3 (root translation torch.Size([55]) """ def __init__(self, tensor_backend, skeleton_tree, is_local): self._skeleton_tree = skeleton_tree self._is_local = is_local self.tensor = tensor_backend.clone() def __len__(self): return self.tensor.shape[0] @property def rotation(self): if not hasattr(self, "_rotation"): self._rotation = self.tensor[..., : self.num_joints * 4].reshape( *(self.tensor.shape[:-1] + (self.num_joints, 4)) ) return self._rotation @property def _local_rotation(self): if self._is_local: return self.rotation else: return None @property def _global_rotation(self): if not self._is_local: return self.rotation else: return None @property def is_local(self): """is the rotation represented in local frame? :rtype: bool """ return self._is_local @property def invariant_property(self): return {"skeleton_tree": self.skeleton_tree, "is_local": self.is_local} @property def num_joints(self): """number of joints in the skeleton tree :rtype: int """ return self.skeleton_tree.num_joints @property def skeleton_tree(self): """skeleton tree :rtype: SkeletonTree """ return self._skeleton_tree @property def root_translation(self): """root translation :rtype: Tensor """ if not hasattr(self, "_root_translation"): self._root_translation = self.tensor[..., self.num_joints * 4 : self.num_joints * 4 + 3] return self._root_translation @property def global_transformation(self): """global transformation of each joint (transform from joint frame to global frame)""" # Forward Kinematics if not hasattr(self, "_global_transformation"): local_transformation = self.local_transformation global_transformation = [] parent_indices = self.skeleton_tree.parent_indices.numpy() # global_transformation = local_transformation.identity_like() for node_index in range(len(self.skeleton_tree)): parent_index = parent_indices[node_index] if parent_index == -1: global_transformation.append(local_transformation[..., node_index, :]) else: global_transformation.append( transform_mul( global_transformation[parent_index], local_transformation[..., node_index, :], ) ) self._global_transformation = torch.stack(global_transformation, axis=-2) return self._global_transformation @property def global_rotation(self): """global rotation of each joint (rotation matrix to rotate from joint's F.O.R to global F.O.R)""" if self._global_rotation is None: if not hasattr(self, "_comp_global_rotation"): self._comp_global_rotation = transform_rotation(self.global_transformation) return self._comp_global_rotation else: return self._global_rotation @property def global_translation(self): """global translation of each joint""" if not hasattr(self, "_global_translation"): self._global_translation = transform_translation(self.global_transformation) return self._global_translation @property def global_translation_xy(self): """global translation in xy""" trans_xy_data = self.global_translation.zeros_like() trans_xy_data[..., 0:2] = self.global_translation[..., 0:2] return trans_xy_data @property def global_translation_xz(self): """global translation in xz""" trans_xz_data = self.global_translation.zeros_like() trans_xz_data[..., 0:1] = self.global_translation[..., 0:1] trans_xz_data[..., 2:3] = self.global_translation[..., 2:3] return trans_xz_data @property def local_rotation(self): """the rotation from child frame to parent frame given in the order of child nodes appeared in `.skeleton_tree.node_names`""" if self._local_rotation is None: if not hasattr(self, "_comp_local_rotation"): local_rotation = quat_identity_like(self.global_rotation) for node_index in range(len(self.skeleton_tree)): parent_index = self.skeleton_tree.parent_indices[node_index] if parent_index == -1: local_rotation[..., node_index, :] = self.global_rotation[ ..., node_index, : ] else: local_rotation[..., node_index, :] = quat_mul_norm( quat_inverse(self.global_rotation[..., parent_index, :]), self.global_rotation[..., node_index, :], ) self._comp_local_rotation = local_rotation return self._comp_local_rotation else: return self._local_rotation @property def local_transformation(self): """local translation + local rotation. It describes the transformation from child frame to parent frame given in the order of child nodes appeared in `.skeleton_tree.node_names` """ if not hasattr(self, "_local_transformation"): self._local_transformation = transform_from_rotation_translation( r=self.local_rotation, t=self.local_translation ) return self._local_transformation @property def local_translation(self): """local translation of the skeleton state. It is identical to the local translation in `.skeleton_tree.local_translation` except the root translation. The root translation is identical to `.root_translation`""" if not hasattr(self, "_local_translation"): broadcast_shape = ( tuple(self.tensor.shape[:-1]) + (len(self.skeleton_tree),) + tuple(self.skeleton_tree.local_translation.shape[-1:]) ) local_translation = self.skeleton_tree.local_translation.broadcast_to( *broadcast_shape ).clone() local_translation[..., 0, :] = self.root_translation self._local_translation = local_translation return self._local_translation # Root Properties @property def root_translation_xy(self): """root translation on xy""" if not hasattr(self, "_root_translation_xy"): self._root_translation_xy = self.global_translation_xy[..., 0, :] return self._root_translation_xy @property def global_root_rotation(self): """root rotation""" if not hasattr(self, "_global_root_rotation"): self._global_root_rotation = self.global_rotation[..., 0, :] return self._global_root_rotation @property def global_root_yaw_rotation(self): """root yaw rotation""" if not hasattr(self, "_global_root_yaw_rotation"): self._global_root_yaw_rotation = self.global_root_rotation.yaw_rotation() return self._global_root_yaw_rotation # Properties relative to root @property def local_translation_to_root(self): """The 3D translation from joint frame to the root frame.""" if not hasattr(self, "_local_translation_to_root"): self._local_translation_to_root = ( self.global_translation - self.root_translation.unsqueeze(-1) ) return self._local_translation_to_root @property def local_rotation_to_root(self): """The 3D rotation from joint frame to the root frame. It is equivalent to The root_R_world * world_R_node""" return quat_inverse(self.global_root_rotation).unsqueeze(-1) * self.global_rotation def compute_forward_vector( self, left_shoulder_index, right_shoulder_index, left_hip_index, right_hip_index, gaussian_filter_width=20, ): """Computes forward vector based on cross product of the up vector with average of the right->left shoulder and hip vectors""" global_positions = self.global_translation # Perpendicular to the forward direction. # Uses the shoulders and hips to find this. side_direction = ( global_positions[:, left_shoulder_index].numpy() - global_positions[:, right_shoulder_index].numpy() + global_positions[:, left_hip_index].numpy() - global_positions[:, right_hip_index].numpy() ) side_direction = side_direction / np.sqrt((side_direction**2).sum(axis=-1))[..., np.newaxis] # Forward direction obtained by crossing with the up direction. forward_direction = np.cross(side_direction, np.array([[0, 1, 0]])) # Smooth the forward direction with a Gaussian. # Axis 0 is the time/frame axis. forward_direction = filters.gaussian_filter1d( forward_direction, gaussian_filter_width, axis=0, mode="nearest" ) forward_direction = ( forward_direction / np.sqrt((forward_direction**2).sum(axis=-1))[..., np.newaxis] ) return torch.from_numpy(forward_direction) @staticmethod def _to_state_vector(rot, rt): state_shape = rot.shape[:-2] vr = rot.reshape(*(state_shape + (-1,))) vt = rt.broadcast_to(*state_shape + rt.shape[-1:]).reshape(*(state_shape + (-1,))) v = torch.cat([vr, vt], axis=-1) return v @classmethod def from_dict( cls: type["SkeletonState"], dict_repr: OrderedDict, *args, **kwargs ) -> "SkeletonState": rot = TensorUtils.from_dict(dict_repr["rotation"], *args, **kwargs) rt = TensorUtils.from_dict(dict_repr["root_translation"], *args, **kwargs) return cls( SkeletonState._to_state_vector(rot, rt), SkeletonTree.from_dict(dict_repr["skeleton_tree"], *args, **kwargs), dict_repr["is_local"], ) def to_dict(self) -> OrderedDict: return OrderedDict( [ ("rotation", tensor_to_dict(self.rotation)), ("root_translation", tensor_to_dict(self.root_translation)), ("skeleton_tree", self.skeleton_tree.to_dict()), ("is_local", self.is_local), ] ) @classmethod def from_rotation_and_root_translation(cls, skeleton_tree, r, t, is_local=True): """ Construct a skeleton state from rotation and root translation :param skeleton_tree: the skeleton tree :type skeleton_tree: SkeletonTree :param r: rotation (either global or local) :type r: Tensor :param t: root translation :type t: Tensor :param is_local: to indicate that whether the rotation is local or global :type is_local: bool, optional, default=True """ assert r.dim() > 0, f"the rotation needs to have at least 1 dimension (dim = {r.dim})" state_vec = SkeletonState._to_state_vector(r, t) return cls( state_vec, skeleton_tree=skeleton_tree, is_local=is_local, ) @classmethod def zero_pose(cls, skeleton_tree): """ Construct a zero-pose skeleton state from the skeleton tree by assuming that all the local rotation is 0 and root translation is also 0. :param skeleton_tree: the skeleton tree as the rigid body :type skeleton_tree: SkeletonTree """ return cls.from_rotation_and_root_translation( skeleton_tree=skeleton_tree, r=quat_identity([skeleton_tree.num_joints]), t=torch.zeros(3, dtype=skeleton_tree.local_translation.dtype), is_local=True, ) def local_repr(self): """ Convert the skeleton state into local representation. This will only affects the values of .tensor. If the skeleton state already has `is_local=True`. This method will do nothing. :rtype: SkeletonState """ if self.is_local: return self return SkeletonState.from_rotation_and_root_translation( self.skeleton_tree, r=self.local_rotation, t=self.root_translation, is_local=True, ) def global_repr(self): """ Convert the skeleton state into global representation. This will only affects the values of .tensor. If the skeleton state already has `is_local=False`. This method will do nothing. :rtype: SkeletonState """ if not self.is_local: return self return SkeletonState.from_rotation_and_root_translation( self.skeleton_tree, r=self.global_rotation, t=self.root_translation, is_local=False, ) def _get_pairwise_average_translation(self): global_transform_inv = transform_inverse(self.global_transformation) p1 = global_transform_inv.unsqueeze(-2) p2 = self.global_transformation.unsqueeze(-3) pairwise_translation = ( transform_translation(transform_mul(p1, p2)) .reshape(-1, len(self.skeleton_tree), len(self.skeleton_tree), 3) .mean(axis=0) ) return pairwise_translation def _transfer_to(self, new_skeleton_tree: SkeletonTree): old_indices = list(map(self.skeleton_tree.index, new_skeleton_tree)) return SkeletonState.from_rotation_and_root_translation( new_skeleton_tree, r=self.global_rotation[..., old_indices, :], t=self.root_translation, is_local=False, ) def drop_nodes_by_names( self, node_names: list[str], estimate_local_translation_from_states: bool = True ) -> "SkeletonState": """ Drop a list of nodes from the skeleton and re-compute the local rotation to match the original joint position as much as possible. :param node_names: a list node names that specifies the nodes need to be dropped :type node_names: List of strings :param estimate_local_translation_from_states: the boolean indicator that specifies whether\ or not to re-estimate the local translation from the states (avg.) :type estimate_local_translation_from_states: boolean :rtype: SkeletonState """ if estimate_local_translation_from_states: pairwise_translation = self._get_pairwise_average_translation() else: pairwise_translation = None new_skeleton_tree = self.skeleton_tree.drop_nodes_by_names(node_names, pairwise_translation) return self._transfer_to(new_skeleton_tree) def keep_nodes_by_names( self, node_names: list[str], estimate_local_translation_from_states: bool = True ) -> "SkeletonState": """ Keep a list of nodes and drop all other nodes from the skeleton and re-compute the local rotation to match the original joint position as much as possible. :param node_names: a list node names that specifies the nodes need to be dropped :type node_names: List of strings :param estimate_local_translation_from_states: the boolean indicator that specifies whether\ or not to re-estimate the local translation from the states (avg.) :type estimate_local_translation_from_states: boolean :rtype: SkeletonState """ return self.drop_nodes_by_names( list(filter(lambda x: (x not in node_names), self)), estimate_local_translation_from_states, ) def _remapped_to(self, joint_mapping: dict[str, str], target_skeleton_tree: SkeletonTree): joint_mapping_inv = {target: source for source, target in joint_mapping.items()} reduced_target_skeleton_tree = target_skeleton_tree.keep_nodes_by_names( list(joint_mapping_inv) ) n_joints = ( len(joint_mapping), len(self.skeleton_tree), len(reduced_target_skeleton_tree), ) assert ( len(set(n_joints)) == 1 ), "the joint mapping is not consistent with the skeleton trees" source_indices = [ self.skeleton_tree.index(joint_mapping_inv[x]) for x in reduced_target_skeleton_tree ] target_local_rotation = self.local_rotation[..., source_indices, :] return SkeletonState.from_rotation_and_root_translation( skeleton_tree=reduced_target_skeleton_tree, r=target_local_rotation, t=self.root_translation, is_local=True, ) def retarget_to( self, joint_mapping: dict[str, str], source_tpose_local_rotation, source_tpose_root_translation: np.ndarray, target_skeleton_tree: SkeletonTree, target_tpose_local_rotation, target_tpose_root_translation: np.ndarray, rotation_to_target_skeleton, scale_to_target_skeleton: float, z_up: bool = True, ) -> "SkeletonState": """ Retarget the skeleton state to a target skeleton tree. This is a naive retarget implementation with rough approximations. The function follows the procedures below. Steps: 1. Drop the joints from the source (self) that do not belong to the joint mapping\ with an implementation that is similar to "keep_nodes_by_names()" - take a\ look at the function doc for more details (same for source_tpose) 2. Rotate the source state and the source tpose by "rotation_to_target_skeleton"\ to align the source with the target orientation 3. Extract the root translation and normalize it to match the scale of the target\ skeleton 4. Extract the global rotation from source state relative to source tpose and\ re-apply the relative rotation to the target tpose to construct the global\ rotation after retargetting 5. Combine the computed global rotation and the root translation from 3 and 4 to\ complete the retargeting. 6. Make feet on the ground (global translation z) :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ the target skeleton :type joint_mapping: dict[str, str] :param source_tpose_local_rotation: the local rotation of the source skeleton :type source_tpose_local_rotation: Tensor :param source_tpose_root_translation: the root translation of the source tpose :type source_tpose_root_translation: np.ndarray :param target_skeleton_tree: the target skeleton tree :type target_skeleton_tree: SkeletonTree :param target_tpose_local_rotation: the local rotation of the target skeleton :type target_tpose_local_rotation: Tensor :param target_tpose_root_translation: the root translation of the target tpose :type target_tpose_root_translation: Tensor :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ the frame of reference of the target skeleton and s is the frame of reference of the source\ skeleton :type rotation_to_target_skeleton: Tensor :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ factor needs to be 0.01. :type scale_to_target_skeleton: float :rtype: SkeletonState """ # STEP 0: Preprocess source_tpose = SkeletonState.from_rotation_and_root_translation( skeleton_tree=self.skeleton_tree, r=source_tpose_local_rotation, t=source_tpose_root_translation, is_local=True, ) target_tpose = SkeletonState.from_rotation_and_root_translation( skeleton_tree=target_skeleton_tree, r=target_tpose_local_rotation, t=target_tpose_root_translation, is_local=True, ) # STEP 1: Drop the irrelevant joints pairwise_translation = self._get_pairwise_average_translation() node_names = list(joint_mapping) new_skeleton_tree = self.skeleton_tree.keep_nodes_by_names(node_names, pairwise_translation) # TODO: combine the following steps before STEP 3 source_tpose = source_tpose._transfer_to(new_skeleton_tree) source_state = self._transfer_to(new_skeleton_tree) source_tpose = source_tpose._remapped_to(joint_mapping, target_skeleton_tree) source_state = source_state._remapped_to(joint_mapping, target_skeleton_tree) # STEP 2: Rotate the source to align with the target new_local_rotation = source_tpose.local_rotation.clone() new_local_rotation[..., 0, :] = quat_mul_norm( rotation_to_target_skeleton, source_tpose.local_rotation[..., 0, :] ) source_tpose = SkeletonState.from_rotation_and_root_translation( skeleton_tree=source_tpose.skeleton_tree, r=new_local_rotation, t=quat_rotate(rotation_to_target_skeleton, source_tpose.root_translation), is_local=True, ) new_local_rotation = source_state.local_rotation.clone() new_local_rotation[..., 0, :] = quat_mul_norm( rotation_to_target_skeleton, source_state.local_rotation[..., 0, :] ) source_state = SkeletonState.from_rotation_and_root_translation( skeleton_tree=source_state.skeleton_tree, r=new_local_rotation, t=quat_rotate(rotation_to_target_skeleton, source_state.root_translation), is_local=True, ) # STEP 3: Normalize to match the target scale root_translation_diff = ( source_state.root_translation - source_tpose.root_translation ) * scale_to_target_skeleton # STEP 4: the global rotation from source state relative to source tpose and # re-apply to the target current_skeleton_tree = source_state.skeleton_tree target_tpose_global_rotation = source_state.global_rotation[0, :].clone() for current_index, name in enumerate(current_skeleton_tree): if name in target_tpose.skeleton_tree: target_tpose_global_rotation[current_index, :] = target_tpose.global_rotation[ target_tpose.skeleton_tree.index(name), : ] global_rotation_diff = quat_mul_norm( source_state.global_rotation, quat_inverse(source_tpose.global_rotation) ) new_global_rotation = quat_mul_norm(global_rotation_diff, target_tpose_global_rotation) # STEP 5: Putting 3 and 4 together current_skeleton_tree = source_state.skeleton_tree shape = source_state.global_rotation.shape[:-1] shape = shape[:-1] + target_tpose.global_rotation.shape[-2:-1] new_global_rotation_output = quat_identity(shape) for current_index, name in enumerate(target_skeleton_tree): while name not in current_skeleton_tree: name = target_skeleton_tree.parent_of(name) parent_index = current_skeleton_tree.index(name) new_global_rotation_output[:, current_index, :] = new_global_rotation[ :, parent_index, : ] source_state = SkeletonState.from_rotation_and_root_translation( skeleton_tree=target_skeleton_tree, r=new_global_rotation_output, t=target_tpose.root_translation + root_translation_diff, is_local=False, ).local_repr() return source_state def retarget_to_by_tpose( self, joint_mapping: dict[str, str], source_tpose: "SkeletonState", target_tpose: "SkeletonState", rotation_to_target_skeleton, scale_to_target_skeleton: float, ) -> "SkeletonState": """ Retarget the skeleton state to a target skeleton tree. This is a naive retarget implementation with rough approximations. See the method `retarget_to()` for more information :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ the target skeleton :type joint_mapping: dict[str, str] :param source_tpose: t-pose of the source skeleton :type source_tpose: SkeletonState :param target_tpose: t-pose of the target skeleton :type target_tpose: SkeletonState :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ the frame of reference of the target skeleton and s is the frame of reference of the source\ skeleton :type rotation_to_target_skeleton: Tensor :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ factor needs to be 0.01. :type scale_to_target_skeleton: float :rtype: SkeletonState """ assert ( len(source_tpose.shape) == 0 and len(target_tpose.shape) == 0 ), "the retargeting script currently doesn't support vectorized operations" return self.retarget_to( joint_mapping, source_tpose.local_rotation, source_tpose.root_translation, target_tpose.skeleton_tree, target_tpose.local_rotation, target_tpose.root_translation, rotation_to_target_skeleton, scale_to_target_skeleton, ) class SkeletonMotion(SkeletonState): def __init__(self, tensor_backend, skeleton_tree, is_local, fps, *args, **kwargs): self._fps = fps super().__init__(tensor_backend, skeleton_tree, is_local, *args, **kwargs) def clone(self): return SkeletonMotion(self.tensor.clone(), self.skeleton_tree, self._is_local, self._fps) @property def invariant_property(self): return { "skeleton_tree": self.skeleton_tree, "is_local": self.is_local, "fps": self.fps, } @property def global_velocity(self): """global velocity""" curr_index = self.num_joints * 4 + 3 return self.tensor[..., curr_index : curr_index + self.num_joints * 3].reshape( *(self.tensor.shape[:-1] + (self.num_joints, 3)) ) @property def global_angular_velocity(self): """global angular velocity""" curr_index = self.num_joints * 7 + 3 return self.tensor[..., curr_index : curr_index + self.num_joints * 3].reshape( *(self.tensor.shape[:-1] + (self.num_joints, 3)) ) @property def fps(self): """number of frames per second""" return self._fps @property def time_delta(self): """time between two adjacent frames""" return 1.0 / self.fps @property def global_root_velocity(self): """global root velocity""" return self.global_velocity[..., 0, :] @property def global_root_angular_velocity(self): """global root angular velocity""" return self.global_angular_velocity[..., 0, :] @classmethod def from_state_vector_and_velocity( cls, skeleton_tree, state_vector, global_velocity, global_angular_velocity, is_local, fps, ): """ Construct a skeleton motion from a skeleton state vector, global velocity and angular velocity at each joint. :param skeleton_tree: the skeleton tree that the motion is based on :type skeleton_tree: SkeletonTree :param state_vector: the state vector from the skeleton state by `.tensor` :type state_vector: Tensor :param global_velocity: the global velocity at each joint :type global_velocity: Tensor :param global_angular_velocity: the global angular velocity at each joint :type global_angular_velocity: Tensor :param is_local: if the rotation ins the state vector is given in local frame :type is_local: boolean :param fps: number of frames per second :type fps: int :rtype: SkeletonMotion """ state_shape = state_vector.shape[:-1] v = global_velocity.reshape(*(state_shape + (-1,))) av = global_angular_velocity.reshape(*(state_shape + (-1,))) new_state_vector = torch.cat([state_vector, v, av], axis=-1) return cls( new_state_vector, skeleton_tree=skeleton_tree, is_local=is_local, fps=fps, ) @classmethod def from_skeleton_state(cls: type["SkeletonMotion"], skeleton_state: SkeletonState, fps: int): """ Construct a skeleton motion from a skeleton state. The velocities are estimated using second order guassian filter along the last axis. The skeleton state must have at least .dim >= 1 :param skeleton_state: the skeleton state that the motion is based on :type skeleton_state: SkeletonState :param fps: number of frames per second :type fps: int :rtype: SkeletonMotion """ assert ( type(skeleton_state) == SkeletonState ), f"expected type of {SkeletonState}, got {type(skeleton_state)}" global_velocity = SkeletonMotion._compute_velocity( p=skeleton_state.global_translation, time_delta=1 / fps ) global_angular_velocity = SkeletonMotion._compute_angular_velocity( r=skeleton_state.global_rotation, time_delta=1 / fps ) return cls.from_state_vector_and_velocity( skeleton_tree=skeleton_state.skeleton_tree, state_vector=skeleton_state.tensor, global_velocity=global_velocity, global_angular_velocity=global_angular_velocity, is_local=skeleton_state.is_local, fps=fps, ) @staticmethod def _to_state_vector(rot, rt, vel, avel): state_shape = rot.shape[:-2] skeleton_state_v = SkeletonState._to_state_vector(rot, rt) v = vel.reshape(*(state_shape + (-1,))) av = avel.reshape(*(state_shape + (-1,))) skeleton_motion_v = torch.cat([skeleton_state_v, v, av], axis=-1) return skeleton_motion_v @classmethod def from_dict( cls: type["SkeletonMotion"], dict_repr: OrderedDict, *args, **kwargs ) -> "SkeletonMotion": rot = TensorUtils.from_dict(dict_repr["rotation"], *args, **kwargs) rt = TensorUtils.from_dict(dict_repr["root_translation"], *args, **kwargs) vel = TensorUtils.from_dict(dict_repr["global_velocity"], *args, **kwargs) avel = TensorUtils.from_dict(dict_repr["global_angular_velocity"], *args, **kwargs) return cls( SkeletonMotion._to_state_vector(rot, rt, vel, avel), skeleton_tree=SkeletonTree.from_dict(dict_repr["skeleton_tree"], *args, **kwargs), is_local=dict_repr["is_local"], fps=dict_repr["fps"], ) def to_dict(self) -> OrderedDict: return OrderedDict( [ ("rotation", tensor_to_dict(self.rotation)), ("root_translation", tensor_to_dict(self.root_translation)), ("global_velocity", tensor_to_dict(self.global_velocity)), ( "global_angular_velocity", tensor_to_dict(self.global_angular_velocity), ), ("skeleton_tree", self.skeleton_tree.to_dict()), ("is_local", self.is_local), ("fps", self.fps), ] ) # @classmethod # def from_fbx( # cls: type["SkeletonMotion"], # fbx_file_path, # fbx_configs, # skeleton_tree=None, # is_local=True, # fps=120, # root_joint="", # root_trans_index=0, # *args, # **kwargs, # ) -> "SkeletonMotion": # """ # Construct a skeleton motion from a fbx file (TODO - generalize this). If the skeleton tree # is not given, it will use the first frame of the mocap to construct the skeleton tree. # :param fbx_file_path: the path of the fbx file # :type fbx_file_path: string # :param fbx_configs: the configuration in terms of {"tmp_path": ..., "fbx_py27_path": ...} # :type fbx_configs: dict # :param skeleton_tree: the optional skeleton tree that the rotation will be applied to # :type skeleton_tree: SkeletonTree, optional # :param is_local: the state vector uses local or global rotation as the representation # :type is_local: bool, optional, default=True # :rtype: SkeletonMotion # """ # joint_names, joint_parents, transforms, fps = fbx_to_array(fbx_file_path, fbx_configs, root_joint, fps) # # swap the last two axis to match the convention # local_transform = euclidean_to_transform(transformation_matrix=torch.from_numpy(np.swapaxes(np.array(transforms), -1, -2),).float()) # local_rotation = transform_rotation(local_transform) # root_translation = transform_translation(local_transform)[..., root_trans_index, :] # joint_parents = torch.from_numpy(np.array(joint_parents)).int() # if skeleton_tree is None: # local_translation = transform_translation(local_transform).reshape(-1, len(joint_parents), 3)[0] # skeleton_tree = SkeletonTree(joint_names, joint_parents, local_translation) # skeleton_state = SkeletonState.from_rotation_and_root_translation(skeleton_tree, r=local_rotation, t=root_translation, is_local=True) # if not is_local: # skeleton_state = skeleton_state.global_repr() # return cls.from_skeleton_state(skeleton_state=skeleton_state, fps=fps) @staticmethod def _compute_velocity(p, time_delta, guassian_filter=True): velocity = np.gradient(p.numpy(), axis=-3) / time_delta if guassian_filter: velocity = torch.from_numpy( filters.gaussian_filter1d(velocity, 2, axis=-3, mode="nearest") ).to(p) else: velocity = torch.from_numpy(velocity).to(p) return velocity @staticmethod def _compute_angular_velocity(r, time_delta: float, guassian_filter=True): # assume the second last dimension is the time axis diff_quat_data = quat_identity_like(r).to(r) diff_quat_data[..., :-1, :, :] = quat_mul_norm( r[..., 1:, :, :], quat_inverse(r[..., :-1, :, :]) ) diff_angle, diff_axis = quat_angle_axis(diff_quat_data) angular_velocity = diff_axis * diff_angle.unsqueeze(-1) / time_delta if guassian_filter: angular_velocity = torch.from_numpy( filters.gaussian_filter1d(angular_velocity.numpy(), 2, axis=-3, mode="nearest"), ) return angular_velocity def crop(self, start: int, end: int, fps: int | None = None): """ Crop the motion along its last axis. This is equivalent to performing a slicing on the object with [..., start: end: skip_every] where skip_every = old_fps / fps. Note that the new fps provided must be a factor of the original fps. :param start: the beginning frame index :type start: int :param end: the ending frame index :type end: int :param fps: number of frames per second in the output (if not given the original fps will be used) :type fps: int, optional :rtype: SkeletonMotion """ if fps is None: new_fps = int(self.fps) old_fps = int(self.fps) else: new_fps = int(fps) old_fps = int(self.fps) assert old_fps % fps == 0, ( "the resampling doesn't support fps with non-integer division " f"from the original fps: {old_fps} => {fps}" ) skip_every = old_fps // new_fps s = slice(start, end, skip_every) z = self[..., s] rot = z.local_rotation if z.is_local else z.global_rotation rt = z.root_translation vel = z.global_velocity avel = z.global_angular_velocity return SkeletonMotion( SkeletonMotion._to_state_vector(rot, rt, vel, avel), skeleton_tree=z.skeleton_tree, is_local=z.is_local, fps=new_fps, ) def retarget_to( self, joint_mapping: dict[str, str], source_tpose_local_rotation, source_tpose_root_translation: np.ndarray, target_skeleton_tree: "SkeletonTree", target_tpose_local_rotation, target_tpose_root_translation: np.ndarray, rotation_to_target_skeleton, scale_to_target_skeleton: float, z_up: bool = True, ) -> "SkeletonMotion": """ Same as the one in :class:`SkeletonState`. This method discards all velocity information before retargeting and re-estimate the velocity after the retargeting. The same fps is used in the new retargetted motion. :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ the target skeleton :type joint_mapping: dict[str, str] :param source_tpose_local_rotation: the local rotation of the source skeleton :type source_tpose_local_rotation: Tensor :param source_tpose_root_translation: the root translation of the source tpose :type source_tpose_root_translation: np.ndarray :param target_skeleton_tree: the target skeleton tree :type target_skeleton_tree: SkeletonTree :param target_tpose_local_rotation: the local rotation of the target skeleton :type target_tpose_local_rotation: Tensor :param target_tpose_root_translation: the root translation of the target tpose :type target_tpose_root_translation: Tensor :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ the frame of reference of the target skeleton and s is the frame of reference of the source\ skeleton :type rotation_to_target_skeleton: Tensor :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ factor needs to be 0.01. :type scale_to_target_skeleton: float :rtype: SkeletonMotion """ return SkeletonMotion.from_skeleton_state( super().retarget_to( joint_mapping, source_tpose_local_rotation, source_tpose_root_translation, target_skeleton_tree, target_tpose_local_rotation, target_tpose_root_translation, rotation_to_target_skeleton, scale_to_target_skeleton, z_up, ), self.fps, ) def retarget_to_by_tpose( self, joint_mapping: dict[str, str], source_tpose: "SkeletonState", target_tpose: "SkeletonState", rotation_to_target_skeleton, scale_to_target_skeleton: float, z_up: bool = True, ) -> "SkeletonMotion": """ Same as the one in :class:`SkeletonState`. This method discards all velocity information before retargeting and re-estimate the velocity after the retargeting. The same fps is used in the new retargetted motion. :param joint_mapping: a dictionary of that maps the joint node from the source skeleton to \ the target skeleton :type joint_mapping: dict[str, str] :param source_tpose: t-pose of the source skeleton :type source_tpose: SkeletonState :param target_tpose: t-pose of the target skeleton :type target_tpose: SkeletonState :param rotation_to_target_skeleton: the rotation that needs to be applied to the source\ skeleton to align with the target skeleton. Essentially the rotation is t_R_s, where t is\ the frame of reference of the target skeleton and s is the frame of reference of the source\ skeleton :type rotation_to_target_skeleton: Tensor :param scale_to_target_skeleton: the factor that needs to be multiplied from source\ skeleton to target skeleton (unit in distance). For example, to go from `cm` to `m`, the \ factor needs to be 0.01. :type scale_to_target_skeleton: float :rtype: SkeletonMotion """ return self.retarget_to( joint_mapping, source_tpose.local_rotation, source_tpose.root_translation, target_tpose.skeleton_tree, target_tpose.local_rotation, target_tpose.root_translation, rotation_to_target_skeleton, scale_to_target_skeleton, z_up, )