diff --git a/newton/examples/example_selection_ant.py b/newton/examples/example_selection_ant.py index 3d79d0e630..14ce765370 100644 --- a/newton/examples/example_selection_ant.py +++ b/newton/examples/example_selection_ant.py @@ -76,7 +76,7 @@ def __init__(self, stage_path=None, num_envs=8): # =========================================================== # create articulation view # =========================================================== - self.ants = ArticulationView(self.model, "/World/envs/*/Robot/torso", include_free_joint=True) + self.ants = ArticulationView(self.model, "/World/envs/*/Robot/torso") print(f"articulation count: {self.ants.count}") print(f"link_count: {self.ants.link_count}") @@ -95,27 +95,14 @@ def __init__(self, stage_path=None, num_envs=8): dof_limit_upper = wp.to_torch(self.ants.get_attribute("joint_limit_upper", self.model)) default_dof_transforms = 0.5 * (dof_limit_lower + dof_limit_upper) - if self.ants.include_free_joint: - # combined root and dof transforms - self.default_transforms = wp.to_torch(self.ants.get_attribute("joint_q", self.model)).clone() - self.default_transforms[:, 2] = 0.8 # z-coordinate of articulation root - self.default_transforms[:, 7:] = default_dof_transforms - # combined root and dof velocities - self.default_velocities = wp.to_torch(self.ants.get_attribute("joint_qd", self.model)).clone() - self.default_velocities[:, 2] = 0.5 * math.pi # rotate about z-axis - self.default_velocities[:, 5] = 5.0 # move up z-axis - else: - # root transforms - self.default_root_transforms = wp.to_torch(self.ants.get_root_transforms(self.model)).clone() - self.default_root_transforms[:, 2] = 0.8 - # dof transforms - self.default_dof_transforms = default_dof_transforms - # root velocities - self.default_root_velocities = wp.to_torch(self.ants.get_root_velocities(self.model)).clone() - self.default_root_velocities[:, 2] = 0.5 * math.pi # rotate about z-axis - self.default_root_velocities[:, 5] = 5.0 # move up z-axis - # dof velocities - self.default_dof_velocities = wp.to_torch(self.ants.get_attribute("joint_qd", self.model)).clone() + # combined root and dof transforms + self.default_transforms = wp.to_torch(self.ants.get_attribute("joint_q", self.model)).clone() + self.default_transforms[:, 2] = 0.8 # z-coordinate of articulation root + self.default_transforms[:, 7:] = default_dof_transforms + # combined root and dof velocities + self.default_velocities = wp.to_torch(self.ants.get_attribute("joint_qd", self.model)).clone() + self.default_velocities[:, 2] = 0.5 * math.pi # rotate about z-axis + self.default_velocities[:, 5] = 5.0 # move up z-axis # create disjoint index groups to alternate between all_indices = torch.arange(num_envs, dtype=torch.int32) @@ -152,10 +139,9 @@ def step(self): # ========================= # apply random controls # ========================= - joint_forces = 300.0 - 600.0 * torch.rand((self.num_envs, 8)) - if self.ants.include_free_joint: - # include the leading root joint (pad with zeros) - joint_forces = torch.cat([torch.zeros((self.num_envs, 6)), joint_forces], axis=1) + joint_forces = wp.to_torch(self.ants.get_attribute("joint_f", self.control)) + # skip the free joint + joint_forces[:, 6:] = 300.0 - 600.0 * torch.rand((self.num_envs, self.ants.joint_axis_count)) self.ants.set_attribute("joint_f", self.control, joint_forces) with wp.ScopedTimer("step", active=False): @@ -169,18 +155,8 @@ def reset(self, indices=None): # ============================== # set transforms and velocities # ============================== - if self.ants.include_free_joint: - # set root and dof transforms together - self.ants.set_attribute("joint_q", self.state_0, self.default_transforms, indices=indices) - # set root and dof velocities together - self.ants.set_attribute("joint_qd", self.state_0, self.default_velocities, indices=indices) - else: - # set root and dof transforms separately - self.ants.set_root_transforms(self.state_0, self.default_root_transforms, indices=indices) - self.ants.set_attribute("joint_q", self.state_0, self.default_dof_transforms, indices=indices) - # set root and dof velocities separately - self.ants.set_root_velocities(self.state_0, self.default_root_velocities, indices=indices) - self.ants.set_attribute("joint_qd", self.state_0, self.default_dof_velocities, indices=indices) + self.ants.set_attribute("joint_q", self.state_0, self.default_transforms, indices=indices) + self.ants.set_attribute("joint_qd", self.state_0, self.default_velocities, indices=indices) if not isinstance(self.solver, newton.solvers.MuJoCoSolver): self.ants.eval_fk(self.state_0, indices=indices) diff --git a/newton/examples/example_selection_humanoid.py b/newton/examples/example_selection_humanoid.py index d84efae1f7..e70ac48156 100644 --- a/newton/examples/example_selection_humanoid.py +++ b/newton/examples/example_selection_humanoid.py @@ -76,7 +76,7 @@ def __init__(self, stage_path=None, num_envs=8): # =========================================================== # create articulation view # =========================================================== - self.humanoids = ArticulationView(self.model, "/World/envs/*/Robot/torso", include_free_joint=False) + self.humanoids = ArticulationView(self.model, "/World/envs/*/Robot/torso") print(f"articulation count: {self.humanoids.count}") print(f"link_count: {self.humanoids.link_count}") @@ -97,28 +97,14 @@ def __init__(self, stage_path=None, num_envs=8): # dof_limit_upper = wp.to_torch(self.humanoids.get_attribute("joint_limit_upper", self.model)) # default_dof_transforms = 0.5 * (dof_limit_lower + dof_limit_upper) - if self.humanoids.include_free_joint: - # combined root and dof transforms - self.default_transforms = wp.to_torch(self.humanoids.get_attribute("joint_q", self.model)).clone() - self.default_transforms[:, 2] = 1.5 # z-coordinate of articulation root - # self.default_transforms[:, 7:] = default_dof_transforms - # combined root and dof velocities - self.default_velocities = wp.to_torch(self.humanoids.get_attribute("joint_qd", self.model)).clone() - # self.default_velocities[:, 2] = 1.0 * math.pi # rotate about z-axis - # self.default_velocities[:, 5] = 5.0 # move up z-axis - else: - # root transforms - self.default_root_transforms = wp.to_torch(self.humanoids.get_root_transforms(self.model)).clone() - self.default_root_transforms[:, 2] = 1.5 - # dof transforms - # self.default_dof_transforms = default_dof_transforms - self.default_dof_transforms = wp.to_torch(self.humanoids.get_attribute("joint_q", self.model)).clone() - # root velocities - self.default_root_velocities = wp.to_torch(self.humanoids.get_root_velocities(self.model)).clone() - # self.default_root_velocities[:, 2] = 1.0 * math.pi # rotate about z-axis - # self.default_root_velocities[:, 5] = 5.0 # move up z-axis - # dof velocities - self.default_dof_velocities = wp.to_torch(self.humanoids.get_attribute("joint_qd", self.model)).clone() + # combined root and dof transforms + self.default_transforms = wp.to_torch(self.humanoids.get_attribute("joint_q", self.model)).clone() + self.default_transforms[:, 2] = 1.5 # z-coordinate of articulation root + # self.default_transforms[:, 7:] = default_dof_transforms + # combined root and dof velocities + self.default_velocities = wp.to_torch(self.humanoids.get_attribute("joint_qd", self.model)).clone() + # self.default_velocities[:, 2] = 1.0 * math.pi # rotate about z-axis + # self.default_velocities[:, 5] = 5.0 # move up z-axis # create disjoint index groups to alternate between all_indices = torch.arange(num_envs, dtype=torch.int32) @@ -155,10 +141,9 @@ def step(self): # ========================= # apply random controls # ========================= - joint_forces = 20.0 - 40.0 * torch.rand((self.num_envs, self.humanoids.joint_dof_count)) - if self.humanoids.include_free_joint: - # include the leading root joint (pad with zeros) - joint_forces = torch.cat([torch.zeros((self.num_envs, 6)), joint_forces], axis=1) + joint_forces = wp.to_torch(self.humanoids.get_attribute("joint_f", self.control)) + # skip the free joint + joint_forces[:, 6:] = 20.0 - 40.0 * torch.rand((self.num_envs, self.humanoids.joint_axis_count)) self.humanoids.set_attribute("joint_f", self.control, joint_forces) with wp.ScopedTimer("step", active=False): @@ -172,18 +157,8 @@ def reset(self, indices=None): # ============================== # set transforms and velocities # ============================== - if self.humanoids.include_free_joint: - # set root and dof transforms together - self.humanoids.set_attribute("joint_q", self.state_0, self.default_transforms, indices=indices) - # set root and dof velocities together - self.humanoids.set_attribute("joint_qd", self.state_0, self.default_velocities, indices=indices) - else: - # set root and dof transforms separately - self.humanoids.set_root_transforms(self.state_0, self.default_root_transforms, indices=indices) - self.humanoids.set_attribute("joint_q", self.state_0, self.default_dof_transforms, indices=indices) - # set root and dof velocities separately - self.humanoids.set_root_velocities(self.state_0, self.default_root_velocities, indices=indices) - self.humanoids.set_attribute("joint_qd", self.state_0, self.default_dof_velocities, indices=indices) + self.humanoids.set_attribute("joint_q", self.state_0, self.default_transforms, indices=indices) + self.humanoids.set_attribute("joint_qd", self.state_0, self.default_velocities, indices=indices) if not isinstance(self.solver, newton.solvers.MuJoCoSolver): self.humanoids.eval_fk(self.state_0, indices=indices) diff --git a/newton/utils/selection.py b/newton/utils/selection.py index 33226351f1..6e17e86380 100644 --- a/newton/utils/selection.py +++ b/newton/utils/selection.py @@ -95,106 +95,10 @@ def set_mask_indexed_kernel( mask[indices[indices_indices[tid]]] = True -@wp.kernel -def set_articulation_root_transforms_kernel( - articulation_indices: wp.array(dtype=int), - articulation_start: wp.array(dtype=int), - joint_type: wp.array(dtype=int), - joint_q_start: wp.array(dtype=int), - root_transforms: wp.array(dtype=wp.transform), - env_indices: wp.array(dtype=int), - # outputs - joint_q: wp.array(dtype=float), - joint_X_p: wp.array(dtype=wp.transform), -): - tid = wp.tid() - idx = env_indices[tid] - root_pose = root_transforms[idx] - articulation = articulation_indices[idx] - joint_start = articulation_start[articulation] - q_start = joint_q_start[joint_start] - - if joint_type[joint_start] == newton.JOINT_FREE: - for i in range(7): - joint_q[q_start + i] = root_pose[i] - elif joint_type[joint_start] == newton.JOINT_FIXED: - joint_X_p[joint_start] = root_pose - - -@wp.kernel -def get_articulation_root_transforms_kernel( - articulation_indices: wp.array(dtype=int), - articulation_start: wp.array(dtype=int), - joint_parent: wp.array(dtype=int), - joint_child: wp.array(dtype=int), - body_q: wp.array(dtype=wp.transform), - # outputs - root_xforms: wp.array(dtype=wp.transform), -): - tid = wp.tid() - articulation = articulation_indices[tid] - joint_start = articulation_start[articulation] - - if joint_parent[joint_start] != -1: - root_body = joint_parent[joint_start] - else: - root_body = joint_child[joint_start] - - root_pose = body_q[root_body] - - root_xforms[tid] = root_pose - - -@wp.kernel -def set_articulation_root_velocities_kernel( - articulation_indices: wp.array(dtype=int), - articulation_start: wp.array(dtype=int), - joint_type: wp.array(dtype=int), - joint_qd_start: wp.array(dtype=int), - root_vels: wp.array(dtype=wp.spatial_vector), - env_indices: wp.array(dtype=int), - # outputs - joint_qd: wp.array(dtype=float), -): - tid = wp.tid() - idx = env_indices[tid] - articulation = articulation_indices[idx] - joint_start = articulation_start[articulation] - qd_start = joint_qd_start[joint_start] - root_vel = root_vels[idx] - - if joint_type[joint_start] == newton.JOINT_FREE: - for i in range(6): - joint_qd[qd_start + i] = root_vel[i] - - -@wp.kernel -def get_articulation_root_velocities_kernel( - articulation_indices: wp.array(dtype=int), - articulation_start: wp.array(dtype=int), - joint_parent: wp.array(dtype=int), - joint_child: wp.array(dtype=int), - body_qd: wp.array(dtype=wp.spatial_vector), - # outputs - root_vels: wp.array(dtype=wp.spatial_vector), -): - tid = wp.tid() - articulation = articulation_indices[tid] - joint_start = articulation_start[articulation] - - if joint_parent[joint_start] != -1: - root_body = joint_parent[joint_start] - else: - root_body = joint_child[joint_start] - - root_vels[tid] = body_qd[root_body] - - class ArticulationView: - def __init__(self, model: Model, pattern: str, include_free_joint: bool = False, verbose: bool | None = None): + def __init__(self, model: Model, pattern: str, verbose: bool | None = None): self.model = model self.device = model.device - self.include_free_joint = include_free_joint if verbose is None: verbose = wp.config.verbose @@ -210,7 +114,6 @@ def __init__(self, model: Model, pattern: str, include_free_joint: bool = False, # FIXME: avoid this readback? articulation_start = model.articulation_start.numpy() - joint_type = model.joint_type.numpy() joint_parent = model.joint_parent.numpy() joint_child = model.joint_child.numpy() joint_axis_start = model.joint_axis_start.numpy() @@ -251,10 +154,6 @@ def __init__(self, model: Model, pattern: str, include_free_joint: bool = False, for body_id in links: print(f" {model.body_key[body_id]}") - # if the root joint is a free joint, skip it - if joint_type[joint_begin] == newton.JOINT_FREE and not include_free_joint: - joint_begin += 1 - joint_coord_begin = joint_q_start[joint_begin] joint_coord_end = joint_q_start[joint_end] joint_dof_begin = joint_qd_start[joint_begin] @@ -328,155 +227,6 @@ def set_attribute(self, name: str, target: Model | State | Control, values, indi values = wp.indexedarray(values, [indices]) wp.copy(attrib, values) - # convenience wrappers to align with legacy tensor API - # TODO: do we want this? - - # def get_link_transforms(self, source, copy=False): - # return self.get_attribute("body_q", source, copy=copy) - - # def get_link_velocities(self, source, copy=False): - # return self.get_attribute("body_qd", source, copy=copy) - - # ... - - def get_root_transforms(self, source: Model | State): - """ - Get the root transforms of the articulations. - - Args: - source (Model | State): Where to get the root transforms (Model or State). - - Returns: - array: The root transforms (dtype=wp.transform). - """ - if self._root_transforms is None: - self._root_transforms = wp.empty(self.count, dtype=wp.transform, device=self.device) - - wp.launch( - get_articulation_root_transforms_kernel, - self.count, - inputs=[ - self.articulation_indices, - self.model.articulation_start, - self.model.joint_parent, - self.model.joint_child, - source.body_q, - ], - outputs=[ - self._root_transforms, - ], - device=self.device, - ) - - return self._root_transforms - - def set_root_transforms(self, target: Model | State, root_transforms: wp.array, indices=None): - """ - Set the root transforms of the articulations. - Call `eval_fk()` to apply changes to all articulation links. - - Args: - target (Model | State): Where to set the root transforms (Model or State). - root_transforms (array): The root transforms to set (dtype=wp.transform). - """ - - if not is_array(root_transforms): - root_transforms = wp.array(root_transforms, dtype=wp.transform, device=self.device) - - assert len(root_transforms) == self.count, "Root poses should be provided for each articulation" - - if indices is not None: - if not is_array(indices): - indices = wp.array(indices, dtype=int, device=self.device) - else: - indices = self.all_indices - - wp.launch( - set_articulation_root_transforms_kernel, - indices.size, - inputs=[ - self.articulation_indices, - self.model.articulation_start, - self.model.joint_type, - self.model.joint_q_start, - root_transforms, - indices, - ], - outputs=[ - target.joint_q, - self.model.joint_X_p, # hmmm - ], - device=self.device, - ) - - def get_root_velocities(self, source: Model | State): - """ - Get the root velocities of the articulations. - - Args: - source (Model | State): Where to get the root velocities (Model or State). - - Returns: - array: The root velocities (dtype=wp.spatial_vector). - """ - if self._root_velocities is None: - self._root_velocities = wp.empty(self.count, dtype=wp.spatial_vector, device=self.device) - - wp.launch( - get_articulation_root_velocities_kernel, - self.count, - inputs=[ - self.articulation_indices, - self.model.articulation_start, - self.model.joint_parent, - self.model.joint_child, - source.body_qd, - ], - outputs=[ - self._root_velocities, - ], - device=self.device, - ) - - return self._root_velocities - - def set_root_velocities(self, target: Model | State, root_vels: wp.array, indices=None): - """ - Set the root velocities of the articulations. - - Args: - target (Model | State): Where to set the root velocities (Model or State). - root_vels (array): The root velocities to set (dtype=wp.spatial_vector). - """ - - if not is_array(root_vels): - root_vels = wp.array(root_vels, dtype=wp.spatial_vector, device=self.device) - - assert len(root_vels) == self.count, "Root velocities should be provided for each articulation" - - if indices is not None: - if not is_array(indices): - indices = wp.array(indices, dtype=int, device=self.device) - else: - indices = self.all_indices - - wp.launch( - set_articulation_root_velocities_kernel, - indices.size, - inputs=[ - self.articulation_indices, - self.model.articulation_start, - self.model.joint_type, - self.model.joint_qd_start, - root_vels, - indices, - ], - outputs=[ - target.joint_qd, - ], - device=self.device, - ) - def eval_fk(self, target: Model | State, indices=None): if indices is not None: # create a custom mask for builtin eval_fk()