From 6a4f8eb19953bc22e2102df8177cc7012c3cc2bb Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 2 Jul 2025 16:13:40 -0400 Subject: [PATCH] Selection API first draft Signed-off-by: Lukasz Wawrzyniak Ruff --- .../example_selection_articulations.py | 338 ++++++++++ newton/examples/example_selection_cartpole.py | 213 ++++++ newton/sim/model.py | 65 ++ newton/utils/isaaclab.py | 113 ++++ newton/utils/selection.py | 623 ++++++++++++++++++ 5 files changed, 1352 insertions(+) create mode 100644 newton/examples/example_selection_articulations.py create mode 100644 newton/examples/example_selection_cartpole.py create mode 100644 newton/utils/isaaclab.py create mode 100644 newton/utils/selection.py diff --git a/newton/examples/example_selection_articulations.py b/newton/examples/example_selection_articulations.py new file mode 100644 index 0000000000..ce5e333e2c --- /dev/null +++ b/newton/examples/example_selection_articulations.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warp as wp + +import newton +import newton.examples +import newton.utils +from newton.examples import compute_env_offsets +from newton.utils.selection import ArticulationView + +USE_TORCH = False +COLLAPSE_FIXED_JOINTS = False +VERBOSE = True + + +@wp.kernel +def compute_middle_kernel( + lower: wp.array2d(dtype=float), upper: wp.array2d(dtype=float), middle: wp.array2d(dtype=float) +): + i, j = wp.tid() + middle[i, j] = 0.5 * (lower[i, j] + upper[i, j]) + + +@wp.kernel +def init_masks(mask_0: wp.array(dtype=bool), mask_1: wp.array(dtype=bool)): + tid = wp.tid() + yes = tid % 2 == 0 + mask_0[tid] = yes + mask_1[tid] = not yes + + +@wp.kernel +def reset_kernel( + ant_root_velocities: wp.array(dtype=wp.spatial_vector), + hum_root_velocities: wp.array(dtype=wp.spatial_vector), + mask: wp.array(dtype=bool), # optional, can be None + seed: int, +): + tid = wp.tid() + + if mask: + do_it = mask[tid] + else: + do_it = True + + if do_it: + rng = wp.rand_init(seed, tid) + spin_vel = 4.0 * wp.pi * (0.5 - wp.randf(rng)) + jump_vel = 3.0 * wp.randf(rng) + ant_root_velocities[tid] = wp.spatial_vector(0.0, 0.0, spin_vel, 0.0, 0.0, jump_vel) + hum_root_velocities[tid] = wp.spatial_vector(0.0, 0.0, -spin_vel, 0.0, 0.0, jump_vel) + + +@wp.kernel +def random_forces_kernel(dof_forces: wp.array2d(dtype=float), seed: int, num_envs: int): + i, j = wp.tid() + rng = wp.rand_init(seed, i * num_envs + j) + dof_forces[i, j] = 5.0 - 10.0 * wp.randf(rng) + + +class Example: + def __init__(self, stage_path=None, num_envs=8): + self.num_envs = num_envs + + up_axis = newton.Axis.Z + + env_builder = newton.ModelBuilder(up_axis=up_axis) + newton.utils.parse_mjcf( + newton.examples.get_asset("nv_ant.xml"), + env_builder, + ignore_names=["floor", "ground"], + up_axis=up_axis, + xform=wp.transform((0.0, 0.0, 1.0), wp.quat_identity()), + collapse_fixed_joints=COLLAPSE_FIXED_JOINTS, + ) + newton.utils.parse_mjcf( + newton.examples.get_asset("nv_humanoid.xml"), + env_builder, + ignore_names=["floor", "ground"], + up_axis=up_axis, + xform=wp.transform((0.0, 0.0, 3.5), wp.quat_identity()), + collapse_fixed_joints=COLLAPSE_FIXED_JOINTS, + ) + + env_offsets = compute_env_offsets(num_envs, env_offset=(4.0, 4.0, 0.0), up_axis=up_axis) + + builder = newton.ModelBuilder() + for i in range(self.num_envs): + builder.add_builder(env_builder, xform=wp.transform(env_offsets[i], wp.quat_identity())) + + builder.add_ground_plane() + + # finalize model + self.model = builder.finalize() + + self.solver = newton.solvers.MuJoCoSolver(self.model) + + self.renderer = None + if stage_path: + self.renderer = newton.utils.SimRendererOpenGL( + path=stage_path, + model=self.model, + scaling=2.0, + up_axis=str(up_axis), + screen_width=1280, + screen_height=720, + camera_pos=(0, 4, 30), + ) + + self.state_0 = self.model.state() + self.state_1 = self.model.state() + self.control = self.model.control() + + self.sim_time = 0.0 + fps = 60 + self.frame_dt = 1.0 / fps + + self.sim_substeps = 10 + self.sim_dt = self.frame_dt / self.sim_substeps + + self.next_reset = 0.0 + self.step_count = 0 + + # =========================================================== + # create articulation views + # =========================================================== + self.ants = ArticulationView(self.model, "ant", verbose=VERBOSE, exclude_joint_types=[newton.JOINT_FREE]) + self.hums = ArticulationView(self.model, "humanoid", verbose=VERBOSE, exclude_joint_types=[newton.JOINT_FREE]) + + if USE_TORCH: + import torch # noqa: PLC0415 + + # default ant root states + self.default_ant_root_transforms = wp.to_torch(self.ants.get_root_transforms(self.model)).clone() + self.default_ant_root_velocities = wp.to_torch(self.ants.get_root_velocities(self.model)).clone() + + # set ant DOFs to the middle of their range by default + dof_limit_lower = wp.to_torch(self.ants.get_attribute("joint_limit_lower", self.model)) + dof_limit_upper = wp.to_torch(self.ants.get_attribute("joint_limit_upper", self.model)) + self.default_ant_dof_positions = 0.5 * (dof_limit_lower + dof_limit_upper) + self.default_ant_dof_velocities = wp.to_torch(self.ants.get_dof_velocities(self.model)).clone() + + # default humanoid states + self.default_hum_root_transforms = wp.to_torch(self.hums.get_root_transforms(self.model)).clone() + self.default_hum_root_velocities = wp.to_torch(self.hums.get_root_velocities(self.model)).clone() + self.default_hum_dof_positions = wp.to_torch(self.hums.get_dof_positions(self.model)).clone() + self.default_hum_dof_velocities = wp.to_torch(self.hums.get_dof_velocities(self.model)).clone() + + # create disjoint subsets to alternate resets + all_indices = torch.arange(num_envs, dtype=torch.int32) + self.mask_0 = torch.zeros(num_envs, dtype=bool) + self.mask_0[all_indices[::2]] = True + self.mask_1 = torch.zeros(num_envs, dtype=bool) + self.mask_1[all_indices[1::2]] = True + else: + # default ant root states + self.default_ant_root_transforms = wp.clone(self.ants.get_root_transforms(self.model)) + self.default_ant_root_velocities = wp.clone(self.ants.get_root_velocities(self.model)) + + # set ant DOFs to the middle of their range by default + dof_limit_lower = self.ants.get_attribute("joint_limit_lower", self.model) + dof_limit_upper = self.ants.get_attribute("joint_limit_upper", self.model) + self.default_ant_dof_positions = wp.empty_like(dof_limit_lower) + wp.launch( + compute_middle_kernel, + dim=self.default_ant_dof_positions.shape, + inputs=[dof_limit_lower, dof_limit_upper, self.default_ant_dof_positions], + ) + self.default_ant_dof_velocities = wp.clone(self.ants.get_dof_velocities(self.model)) + + # default humanoid states + self.default_hum_root_transforms = wp.clone(self.hums.get_root_transforms(self.model)) + self.default_hum_root_velocities = wp.clone(self.hums.get_root_velocities(self.model)) + self.default_hum_dof_positions = wp.clone(self.hums.get_dof_positions(self.model)) + self.default_hum_dof_velocities = wp.clone(self.hums.get_dof_velocities(self.model)) + + # create disjoint subsets to alternate resets + self.mask_0 = wp.empty(num_envs, dtype=bool) + self.mask_1 = wp.empty(num_envs, dtype=bool) + wp.launch(init_masks, dim=num_envs, inputs=[self.mask_0, self.mask_1]) + + # reset all + self.reset() + self.next_reset = self.sim_time + 2.0 + + self.use_cuda_graph = wp.get_device().is_cuda + if self.use_cuda_graph: + with wp.ScopedCapture() as capture: + self.simulate() + self.graph = capture.graph + + def simulate(self): + for _ in range(self.sim_substeps): + self.state_0.clear_forces() + + # explicit collisions needed without MuJoCo solver + if not isinstance(self.solver, newton.solvers.MuJoCoSolver): + contacts = self.model.collide(self.state_0) + else: + contacts = None + + self.solver.step(self.model, self.state_0, self.state_1, self.control, contacts, self.sim_dt) + self.state_0, self.state_1 = self.state_1, self.state_0 + + def step(self): + if self.sim_time >= self.next_reset: + self.reset(mask=self.mask_0) + self.mask_0, self.mask_1 = self.mask_1, self.mask_0 + self.next_reset = self.sim_time + 2.0 + + # ================================ + # apply random controls + # ================================ + if USE_TORCH: + import torch # noqa: PLC0415 + + dof_forces = 5.0 - 10.0 * torch.rand((self.num_envs, self.ants.joint_dof_count)) + else: + dof_forces = self.ants.get_dof_forces(self.control) + wp.launch(random_forces_kernel, dim=dof_forces.shape, inputs=[dof_forces, self.step_count, self.num_envs]) + + self.ants.set_dof_forces(self.control, dof_forces) + + with wp.ScopedTimer("step", active=False): + if self.use_cuda_graph: + wp.capture_launch(self.graph) + else: + self.simulate() + self.sim_time += self.frame_dt + self.step_count += 1 + + def reset(self, mask=None): + # ================================ + # reset transforms and velocities + # ================================ + + if USE_TORCH: + import torch # noqa: PLC0415 + + # randomize ant velocities + self.default_ant_root_velocities[:, 2] = 4.0 * torch.pi * (0.5 - torch.rand(self.num_envs)) + self.default_ant_root_velocities[:, 5] = 3.0 * torch.rand(self.num_envs) + + # humanoids spin in the opposite direction + self.default_hum_root_velocities[:, 2] = -self.default_ant_root_velocities[:, 2] + # humanoids move up at the same speed + self.default_hum_root_velocities[:, 5] = self.default_ant_root_velocities[:, 5] + else: + wp.launch( + reset_kernel, + dim=self.num_envs, + inputs=[self.default_ant_root_velocities, self.default_hum_root_velocities, mask, self.step_count], + ) + + self.ants.set_root_transforms(self.state_0, self.default_ant_root_transforms, mask=mask) + self.ants.set_root_velocities(self.state_0, self.default_ant_root_velocities, mask=mask) + self.ants.set_dof_positions(self.state_0, self.default_ant_dof_positions, mask=mask) + self.ants.set_dof_velocities(self.state_0, self.default_ant_dof_velocities, mask=mask) + + self.hums.set_root_transforms(self.state_0, self.default_hum_root_transforms, mask=mask) + self.hums.set_root_velocities(self.state_0, self.default_hum_root_velocities, mask=mask) + self.hums.set_dof_positions(self.state_0, self.default_hum_dof_positions, mask=mask) + self.hums.set_dof_velocities(self.state_0, self.default_hum_dof_velocities, mask=mask) + + if not isinstance(self.solver, newton.solvers.MuJoCoSolver): + self.ants.eval_fk(self.state_0, mask=mask) + + def render(self): + if self.renderer is None: + return + + with wp.ScopedTimer("render", active=False): + self.renderer.begin_frame(self.sim_time) + self.renderer.render(self.state_0) + self.renderer.end_frame() + + +# scoped device manager for both Warp and Torch +class ScopedDevice: + def __init__(self, device): + self.warp_scoped_device = wp.ScopedDevice(device) + if USE_TORCH: + import torch # noqa: PLC0415 + + self.torch_scoped_device = torch.device(wp.device_to_torch(device)) + + def __enter__(self): + self.warp_scoped_device.__enter__() + if USE_TORCH: + self.torch_scoped_device.__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.warp_scoped_device.__exit__(exc_type, exc_val, exc_tb) + if USE_TORCH: + self.torch_scoped_device.__exit__(exc_type, exc_val, exc_tb) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.") + parser.add_argument( + "--stage_path", + type=lambda x: None if x == "None" else str(x), + default="example_selection_ant.usd", + help="Path to the output USD file.", + ) + parser.add_argument("--num_frames", type=int, default=1200, help="Total number of frames.") + parser.add_argument("--num_envs", type=int, default=16, help="Total number of simulated environments.") + + args = parser.parse_known_args()[0] + + with ScopedDevice(args.device): + example = Example(stage_path=args.stage_path, num_envs=args.num_envs) + + for _ in range(args.num_frames): + example.step() + example.render() + + # import time + # time.sleep(0.2) + + if example.renderer: + example.renderer.save() diff --git a/newton/examples/example_selection_cartpole.py b/newton/examples/example_selection_cartpole.py new file mode 100644 index 0000000000..f73505f6be --- /dev/null +++ b/newton/examples/example_selection_cartpole.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warp as wp + +import newton +import newton.examples +import newton.utils +from newton.examples import compute_env_offsets +from newton.utils.selection import ArticulationView + +USE_TORCH = False +COLLAPSE_FIXED_JOINTS = False +VERBOSE = True + + +@wp.kernel +def randomize_states_kernel(joint_q: wp.array2d(dtype=float), seed: int): + tid = wp.tid() + rng = wp.rand_init(seed, tid) + joint_q[tid, 0] = 2.0 - 4.0 * wp.randf(rng) + joint_q[tid, 1] = wp.pi / 8.0 - wp.pi / 4.0 * wp.randf(rng) + joint_q[tid, 2] = wp.pi / 8.0 - wp.pi / 4.0 * wp.randf(rng) + + +@wp.kernel +def apply_forces_kernel(joint_q: wp.array2d(dtype=float), joint_f: wp.array2d(dtype=float)): + tid = wp.tid() + if joint_q[tid, 0] > 0.0: + joint_f[tid, 0] = -20.0 + else: + joint_f[tid, 0] = 20.0 + + +class Example: + def __init__(self, stage_path=None, num_envs=8): + self.num_envs = num_envs + + up_axis = newton.Axis.Z + + articulation_builder = newton.ModelBuilder(up_axis=up_axis) + newton.utils.parse_urdf( + newton.examples.get_asset("cartpole.urdf"), + articulation_builder, + up_axis=up_axis, + xform=wp.transform((0.0, 0.0, 2.0), wp.quat_identity()), + collapse_fixed_joints=COLLAPSE_FIXED_JOINTS, + enable_self_collisions=False, + floating=False, + ) + + env_offsets = compute_env_offsets(num_envs, env_offset=(4.0, 4.0, 0.0), up_axis=up_axis) + + builder = newton.ModelBuilder() + for i in range(self.num_envs): + builder.add_builder(articulation_builder, xform=wp.transform(env_offsets[i], wp.quat_identity())) + + # finalize model + self.model = builder.finalize() + + self.sim_time = 0.0 + fps = 60 + self.frame_dt = 1.0 / fps + + self.sim_substeps = 10 + self.sim_dt = self.frame_dt / self.sim_substeps + + self.solver = newton.solvers.MuJoCoSolver(self.model, disable_contacts=True) + + self.state_0 = self.model.state() + self.state_1 = self.model.state() + self.control = self.model.control() + + # ======================= + # get cartpole view + # ======================= + self.cartpoles = ArticulationView(self.model, "cartpole", verbose=VERBOSE) + + # ========================= + # randomize initial state + # ========================= + if USE_TORCH: + import torch # noqa: PLC0415 + + cart_positions = 2.0 - 4.0 * torch.rand(num_envs) + pole1_angles = torch.pi / 8.0 - torch.pi / 4.0 * torch.rand(num_envs) + pole2_angles = torch.pi / 8.0 - torch.pi / 4.0 * torch.rand(num_envs) + joint_q = torch.stack([cart_positions, pole1_angles, pole2_angles], dim=1) + else: + joint_q = self.cartpoles.get_attribute("joint_q", self.state_0) + wp.launch(randomize_states_kernel, dim=num_envs, inputs=[joint_q, 42]) + + self.cartpoles.set_attribute("joint_q", self.state_0, joint_q) + + if not isinstance(self.solver, newton.solvers.MuJoCoSolver): + self.cartpoles.eval_fk(self.state_0) + + self.renderer = None + if stage_path: + self.renderer = newton.utils.SimRendererOpenGL( + path=stage_path, + model=self.model, + scaling=1.0, + up_axis=str(up_axis), + screen_width=1280, + screen_height=720, + camera_pos=(0, 3, 10), + ) + + self.use_cuda_graph = wp.get_device().is_cuda + if self.use_cuda_graph: + with wp.ScopedCapture() as capture: + self.simulate() + self.graph = capture.graph + + def simulate(self): + for _ in range(self.sim_substeps): + self.state_0.clear_forces() + self.solver.step(self.model, self.state_0, self.state_1, self.control, None, self.sim_dt) + self.state_0, self.state_1 = self.state_1, self.state_0 + + def step(self): + # ==================================== + # get observations and apply controls + # ==================================== + if USE_TORCH: + import torch # noqa: PLC0415 + + joint_q = wp.to_torch(self.cartpoles.get_attribute("joint_q", self.state_0)) + joint_f = wp.to_torch(self.cartpoles.get_attribute("joint_f", self.control)) + joint_f[:, 0] = torch.where(joint_q[:, 0] > 0, -20, 20) + else: + joint_q = self.cartpoles.get_attribute("joint_q", self.state_0) + joint_f = self.cartpoles.get_attribute("joint_f", self.control) + wp.launch(apply_forces_kernel, dim=joint_f.shape, inputs=[joint_q, joint_f]) + + self.cartpoles.set_attribute("joint_f", self.control, joint_f) + + # simulate + with wp.ScopedTimer("step", active=False): + if self.use_cuda_graph: + wp.capture_launch(self.graph) + else: + self.simulate() + self.sim_time += self.frame_dt + + def render(self): + if self.renderer is None: + return + + with wp.ScopedTimer("render", active=False): + self.renderer.begin_frame(self.sim_time) + self.renderer.render(self.state_0) + self.renderer.end_frame() + + +# scoped device manager for both Warp and Torch +class ScopedDevice: + def __init__(self, device): + self.warp_scoped_device = wp.ScopedDevice(device) + if USE_TORCH: + import torch # noqa: PLC0415 + + self.torch_scoped_device = torch.device(wp.device_to_torch(device)) + + def __enter__(self): + self.warp_scoped_device.__enter__() + if USE_TORCH: + self.torch_scoped_device.__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.warp_scoped_device.__exit__(exc_type, exc_val, exc_tb) + if USE_TORCH: + self.torch_scoped_device.__exit__(exc_type, exc_val, exc_tb) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.") + parser.add_argument( + "--stage_path", + type=lambda x: None if x == "None" else str(x), + default="example_selection_cartpole.usd", + help="Path to the output USD file.", + ) + parser.add_argument("--num_frames", type=int, default=12000, help="Total number of frames.") + parser.add_argument("--num_envs", type=int, default=16, help="Total number of simulated environments.") + + args = parser.parse_known_args()[0] + + with ScopedDevice(args.device): + example = Example(stage_path=args.stage_path, num_envs=args.num_envs) + + for _ in range(args.num_frames): + example.step() + example.render() + + if example.renderer: + example.renderer.save() diff --git a/newton/sim/model.py b/newton/sim/model.py index 1beae51cab..76cd28c22b 100644 --- a/newton/sim/model.py +++ b/newton/sim/model.py @@ -307,6 +307,48 @@ def __init__(self, device: Devicelike | None = None): self.device = wp.get_device(device) """Device on which the Model was allocated.""" + self.attribute_frequency = {} + """Classify each attribute as per body, per joint, per DOF, etc.""" + + # attributes per body + self.attribute_frequency["body_q"] = "body" + self.attribute_frequency["body_qd"] = "body" + self.attribute_frequency["body_com"] = "body" + self.attribute_frequency["body_inertia"] = "body" + self.attribute_frequency["body_inv_inertia"] = "body" + self.attribute_frequency["body_mass"] = "body" + self.attribute_frequency["body_inv_mass"] = "body" + self.attribute_frequency["body_f"] = "body" + + # attributes per joint + self.attribute_frequency["joint_type"] = "joint" + self.attribute_frequency["joint_parent"] = "joint" + self.attribute_frequency["joint_child"] = "joint" + self.attribute_frequency["joint_ancestor"] = "joint" + self.attribute_frequency["joint_X_p"] = "joint" + self.attribute_frequency["joint_X_c"] = "joint" + self.attribute_frequency["joint_dof_dim"] = "joint" + self.attribute_frequency["joint_enabled"] = "joint" + self.attribute_frequency["joint_twist_lower"] = "joint" + self.attribute_frequency["joint_twist_upper"] = "joint" + + # attributes per joint coord + self.attribute_frequency["joint_q"] = "joint_coord" + + # attributes per joint dof + self.attribute_frequency["joint_qd"] = "joint_dof" + self.attribute_frequency["joint_f"] = "joint_dof" + self.attribute_frequency["joint_armature"] = "joint_dof" + self.attribute_frequency["joint_target"] = "joint_dof" + self.attribute_frequency["joint_axis"] = "joint_dof" + self.attribute_frequency["joint_target_ke"] = "joint_dof" + self.attribute_frequency["joint_target_kd"] = "joint_dof" + self.attribute_frequency["joint_dof_mode"] = "joint_dof" + self.attribute_frequency["joint_limit_lower"] = "joint_dof" + self.attribute_frequency["joint_limit_upper"] = "joint_dof" + self.attribute_frequency["joint_limit_ke"] = "joint_dof" + self.attribute_frequency["joint_limit_kd"] = "joint_dof" + def state(self, requires_grad: bool | None = None) -> State: """Returns a state object for the model @@ -431,3 +473,26 @@ def collide( self._collision_pipeline.iterate_mesh_vertices = iterate_mesh_vertices return self._collision_pipeline.collide(self, state) + + def add_attribute(self, name: str, attrib: wp.array, frequency: str): + """Add a custom attribute to the model""" + if hasattr(self, name): + raise AttributeError(f"Attribute '{name}' already exists") + + if not isinstance(attrib, wp.array): + raise AttributeError(f"Attribute '{name}' must be an array, got {type(attrib)}") + if attrib.device != self.device: + raise AttributeError( + f"Attribute '{name}' must be on the same device as the Model, expected {self.device}, got {attrib.device}" + ) + + setattr(self, name, attrib) + + self.attribute_frequency[name] = frequency + + def get_attribute_frequency(self, name): + """Get the frequency of an attribute, e.g., "body", "joint", etc.""" + frequency = self.attribute_frequency.get(name) + if frequency is None: + raise AttributeError(f"Attribute frequency of '{name}' is not known") + return frequency diff --git a/newton/utils/isaaclab.py b/newton/utils/isaaclab.py new file mode 100644 index 0000000000..037b779f1e --- /dev/null +++ b/newton/utils/isaaclab.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +import warp as wp + +import newton +from newton.examples import compute_env_offsets + + +def replicate_environment( + source, + prototype_path: str, + path_pattern: str, + num_envs: int, + env_spacing: tuple[float], + up_axis: newton.AxisType = "Z", + **usd_kwargs, +) -> tuple[newton.ModelBuilder, dict[str:Any]]: + """ + Replicates a prototype USD environment in Newton. + + Args: + source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance. + prototype_path (str): The USD path where the prototype env is defined, e.g., "/World/envs/env_0". + path_pattern (str): The USD path pattern for replicated envs, e.g., "/World/envs/env_{}". + num_envs (int): Number of replicas to create. + env_spacing (tuple[float]): Environment spacing vector. + up_axis (AxisType): The desired up-vector (should match the USD stage). + **usd_kwargs: Keyword arguments to pass to the USD importer (see `newton.utils.parse_usd()`). + + Returns: + (ModelBuilder, dict): The resulting ModelBuilder containing all replicated environments and a dictionary with USD stage information. + """ + + builder = newton.ModelBuilder(up_axis=up_axis) + + # first, load everything except the prototype env + stage_info = newton.utils.parse_usd( + source, + builder, + ignore_paths=[prototype_path], + **usd_kwargs, + ) + + # up_axis sanity check + stage_up_axis = stage_info.get("up_axis") + if isinstance(stage_up_axis, str) and stage_up_axis.upper() != up_axis.upper(): + print(f"WARNING: up_axis '{up_axis}' does not match USD stage up_axis '{stage_up_axis}'") + + # load just the prototype env + prototype_builder = newton.ModelBuilder(up_axis=up_axis) + newton.utils.parse_usd( + source, + prototype_builder, + root_path=prototype_path, + **usd_kwargs, + ) + + env_offsets = compute_env_offsets(num_envs, env_offset=env_spacing, up_axis=up_axis) + + # clone the prototype env with updated paths + for i in range(num_envs): + body_start = builder.body_count + shape_start = builder.shape_count + joint_start = builder.joint_count + articulation_start = builder.articulation_count + + builder.add_builder(prototype_builder, xform=wp.transform(env_offsets[i], wp.quat_identity())) + + if i > 0: + update_paths( + builder, + prototype_path, + path_pattern.format(i), + body_start=body_start, + shape_start=shape_start, + joint_start=joint_start, + articulation_start=articulation_start, + ) + + return builder, stage_info + + +def update_paths( + builder, old_root, new_root, body_start=None, shape_start=None, joint_start=None, articulation_start=None +): + old_len = len(old_root) + if body_start is not None: + for i in range(body_start, builder.body_count): + builder.body_key[i] = f"{new_root}{builder.body_key[i][old_len:]}" + if shape_start is not None: + for i in range(shape_start, builder.shape_count): + builder.shape_key[i] = f"{new_root}{builder.shape_key[i][old_len:]}" + if joint_start is not None: + for i in range(joint_start, builder.joint_count): + builder.joint_key[i] = f"{new_root}{builder.joint_key[i][old_len:]}" + if articulation_start is not None: + for i in range(articulation_start, builder.articulation_count): + builder.articulation_key[i] = f"{new_root}{builder.articulation_key[i][old_len:]}" diff --git a/newton/utils/selection.py b/newton/utils/selection.py new file mode 100644 index 0000000000..6b5ac6e091 --- /dev/null +++ b/newton/utils/selection.py @@ -0,0 +1,623 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools +from fnmatch import fnmatch +from typing import Any + +import warp as wp +from warp.types import is_array + +import newton.sim +from newton import Control, Model, State +from newton.sim import JOINT_DISTANCE, JOINT_FIXED, JOINT_FREE + + +@wp.kernel +def set_model_articulation_mask_kernel( + view_mask: wp.array(dtype=bool), # mask in ArticulationView + view_to_model_map: wp.array(dtype=int), # maps index in ArticulationView to articulation index in Model + articulation_mask: wp.array(dtype=bool), # output: mask of Model articulation indices +): + """ + Set Model articulation mask from a view mask in an ArticulationView. + """ + tid = wp.tid() + if view_mask[tid]: + articulation_mask[view_to_model_map[tid]] = True + + +@wp.kernel +def set_articulation_attribute_1d_kernel( + view_mask: wp.array(dtype=bool), # mask in ArticulationView + values: Any, # 1d array or indexedarray + attrib: Any, # 1d array or indexedarray +): + i = wp.tid() + if view_mask[i]: + attrib[i] = values[i] + + +@wp.kernel +def set_articulation_attribute_2d_kernel( + view_mask: wp.array(dtype=bool), # mask in ArticulationView + values: Any, # 2d array or indexedarray + attrib: Any, # 2d array or indexedarray +): + i, j = wp.tid() + if view_mask[i]: + attrib[i, j] = values[i, j] + + +@wp.kernel +def set_articulation_attribute_3d_kernel( + view_mask: wp.array(dtype=bool), # mask in ArticulationView + values: Any, # 3d array or indexedarray + attrib: Any, # 3d array or indexedarray +): + i, j, k = wp.tid() + if view_mask[i]: + attrib[i, j, k] = values[i, j, k] + + +@wp.kernel +def set_articulation_attribute_4d_kernel( + view_mask: wp.array(dtype=bool), # mask in ArticulationView + values: Any, # 4d array or indexedarray + attrib: Any, # 4d array or indexedarray +): + i, j, k, l = wp.tid() + if view_mask[i]: + attrib[i, j, k, l] = values[i, j, k, l] + + +# explicit overloads to avoid module reloading +for dtype in [float, int, wp.transform, wp.spatial_vector]: + for src_array_type in [wp.array, wp.indexedarray]: + for dst_array_type in [wp.array, wp.indexedarray]: + wp.overload( + set_articulation_attribute_1d_kernel, + {"values": src_array_type(dtype=dtype, ndim=1), "attrib": dst_array_type(dtype=dtype, ndim=1)}, + ) + wp.overload( + set_articulation_attribute_2d_kernel, + {"values": src_array_type(dtype=dtype, ndim=2), "attrib": dst_array_type(dtype=dtype, ndim=2)}, + ) + wp.overload( + set_articulation_attribute_3d_kernel, + {"values": src_array_type(dtype=dtype, ndim=3), "attrib": dst_array_type(dtype=dtype, ndim=3)}, + ) + wp.overload( + set_articulation_attribute_4d_kernel, + {"values": src_array_type(dtype=dtype, ndim=4), "attrib": dst_array_type(dtype=dtype, ndim=4)}, + ) + + +# NOTE: Python slice objects are not hashable in Python < 3.12, so we use this instead. +class Slice: + def __init__(self, start=None, stop=None): + self.start = start + self.stop = stop + + def __hash__(self): + return hash((self.start, self.stop)) + + def __str__(self): + return f"({self.start}, {self.stop})" + + def get(self): + return slice(self.start, self.stop) + + +class ArticulationView: + def __init__( + self, + model: Model, + pattern: str, + include_joints: list[str | int] | None = None, + exclude_joints: list[str | int] | None = None, + include_links: list[str | int] | None = None, + exclude_links: list[str | int] | None = None, + include_joint_types: list[int] | None = None, + exclude_joint_types: list[int] | None = None, + verbose: bool | None = None, + ): + self.model = model + self.device = model.device + + if verbose is None: + verbose = wp.config.verbose + + articulation_ids = [] + for id, key in enumerate(model.articulation_key): + if fnmatch(key, pattern): + articulation_ids.append(id) + + articulation_count = len(articulation_ids) + if articulation_count == 0: + raise KeyError("No matching articulations") + + # FIXME: avoid/reduce this readback? + model_articulation_start = model.articulation_start.numpy() + model_joint_type = model.joint_type.numpy() + model_joint_child = model.joint_child.numpy() + model_joint_q_start = model.joint_q_start.numpy() + model_joint_qd_start = model.joint_qd_start.numpy() + + # FIXME: + # - this assumes homogeneous envs with one selected articulation per env + # - we're going to have problems if there are any bodies or joints in the "global" env + + arti_0 = articulation_ids[0] + + arti_joint_begin = model_articulation_start[arti_0] + arti_joint_end = model_articulation_start[arti_0 + 1] # FIXME: is this always correct? + arti_joint_count = arti_joint_end - arti_joint_begin + arti_link_count = arti_joint_count + + arti_joint_ids = [] + arti_joint_names = [] + arti_joint_types = [] + arti_link_ids = [] + arti_link_names = [] + + def get_name_from_key(key): + return key.split("/")[-1] + + for idx in range(arti_joint_count): + joint_id = arti_joint_begin + idx + arti_joint_ids.append(int(joint_id)) + arti_joint_names.append(get_name_from_key(model.joint_key[joint_id])) + arti_joint_types.append(int(model_joint_type[joint_id])) + link_id = model_joint_child[joint_id] + arti_link_ids.append(int(link_id)) + arti_link_names.append(get_name_from_key(model.body_key[link_id])) + + # create joint inclusion set + if include_joints is None and include_joint_types is None: + joint_include_indices = set(range(arti_joint_count)) + else: + joint_include_indices = set() + if include_joints is not None: + for id in include_joints: + if isinstance(id, str): + for idx, name in enumerate(arti_joint_names): + if fnmatch(name, id): + joint_include_indices.add(idx) + elif isinstance(id, int): + if id >= 0 and id < arti_joint_count: + joint_include_indices.add(id) + else: + raise TypeError(f"Joint ids must be strings or integers, got {id} of type {type(id)}") + if include_joint_types is not None: + for idx in range(arti_joint_count): + if arti_joint_types[idx] in include_joint_types: + joint_include_indices.add(idx) + + # create joint exclusion set + joint_exclude_indices = set() + if exclude_joints is not None: + for id in exclude_joints: + if isinstance(id, str): + for idx, name in enumerate(arti_joint_names): + if fnmatch(name, id): + joint_exclude_indices.add(idx) + elif isinstance(id, int): + if id >= 0 and id < arti_joint_count: + joint_exclude_indices.add(id) + else: + raise TypeError(f"Joint ids must be strings or integers, got {id} of type {type(id)}") + if exclude_joint_types is not None: + for idx in range(arti_joint_count): + if arti_joint_types[idx] in exclude_joint_types: + joint_exclude_indices.add(idx) + + # create link inclusion set + if include_links is None: + link_include_indices = set(range(arti_link_count)) + else: + link_include_indices = set() + if include_links is not None: + for id in include_links: + if isinstance(id, str): + for idx, name in enumerate(arti_link_names): + if fnmatch(name, id): + link_include_indices.add(idx) + elif isinstance(id, int): + if id >= 0 and id < arti_link_count: + link_include_indices.add(id) + else: + raise TypeError(f"Link ids must be strings or integers, got {id} of type {type(id)}") + + # create link exclusion set + link_exclude_indices = set() + if exclude_links is not None: + for id in exclude_links: + if isinstance(id, str): + for idx, name in enumerate(arti_link_names): + if fnmatch(name, id): + link_exclude_indices.add(idx) + elif isinstance(id, int): + if id >= 0 and id < arti_link_count: + link_exclude_indices.add(id) + else: + raise TypeError(f"Link ids must be strings or integers, got {id} of type {type(id)}") + + # compute selected indices + selected_joint_indices = sorted(joint_include_indices - joint_exclude_indices) + selected_link_indices = sorted(link_include_indices - link_exclude_indices) + + selected_joint_ids = [] + selected_joint_dof_ids = [] + selected_joint_coord_ids = [] + selected_link_ids = [] + + self.joint_names = [] + self.joint_dof_names = [] + self.joint_dof_counts = [] + self.joint_coord_names = [] + self.joint_coord_counts = [] + self.body_names = [] + + # populate info for selected joints and dofs + for idx in selected_joint_indices: + # joint + joint_id = arti_joint_ids[idx] + selected_joint_ids.append(joint_id) + joint_name = get_name_from_key(model.joint_key[joint_id]) + self.joint_names.append(joint_name) + # joint dofs + dof_begin = model_joint_qd_start[joint_id] + dof_end = model_joint_qd_start[joint_id + 1] + dof_count = dof_end - dof_begin + if dof_count == 1: + self.joint_dof_names.append(joint_name) + selected_joint_dof_ids.append(int(dof_begin)) + elif dof_count > 1: + for dof in range(dof_count): + self.joint_dof_names.append(f"{joint_name}:{dof}") + selected_joint_dof_ids.append(int(dof_begin + dof)) + # joint coords + coord_begin = model_joint_q_start[joint_id] + coord_end = model_joint_q_start[joint_id + 1] + coord_count = coord_end - coord_begin + if coord_count == 1: + self.joint_coord_names.append(joint_name) + selected_joint_coord_ids.append(int(coord_begin)) + elif coord_count > 1: + for coord in range(coord_count): + self.joint_coord_names.append(f"{joint_name}:{coord}") + selected_joint_coord_ids.append(int(coord_begin + coord)) + + # populate info for selected links + for idx in selected_link_indices: + body_id = arti_link_ids[idx] + selected_link_ids.append(body_id) + self.body_names.append(get_name_from_key(model.body_key[body_id])) + + # selected counts + self.count = articulation_count + self.joint_count = len(selected_joint_ids) + self.joint_dof_count = len(selected_joint_dof_ids) + self.joint_coord_count = len(selected_joint_coord_ids) + self.link_count = len(selected_link_ids) + + # support custom slicing and indexing + self._arti_joint_begin = int(arti_joint_begin) + self._arti_joint_end = int(arti_joint_end) + self._arti_joint_dof_begin = int(model_joint_qd_start[arti_joint_begin]) + self._arti_joint_dof_end = int(model_joint_qd_start[arti_joint_end]) + self._arti_joint_coord_begin = int(model_joint_q_start[arti_joint_begin]) + self._arti_joint_coord_end = int(model_joint_q_start[arti_joint_end]) + + root_joint_type = arti_joint_types[0] + # fixed base means that all linear and angular degrees of freedom are locked at the root + self.is_fixed_base = root_joint_type == JOINT_FIXED + # floating base means that all linear and angular degrees of freedom are unlocked at the root + # (though there might be constraints like distance) + self.is_floating_base = root_joint_type in (JOINT_FREE, JOINT_DISTANCE) + + def is_contiguous_slice(indices): + n = len(indices) + if n > 1: + for i in range(1, n): + if indices[i] != indices[i - 1] + 1: + return False + return True + + self.joints_contiguous = is_contiguous_slice(selected_joint_ids) + self.joint_dofs_contiguous = is_contiguous_slice(selected_joint_dof_ids) + self.joint_coords_contiguous = is_contiguous_slice(selected_joint_coord_ids) + self.links_contiguous = is_contiguous_slice(selected_link_ids) + + # contiguous slices or indices by attribute frequency + # + # FIXME: guard against empty selections + # + self._frequency_slices = {} + self._frequency_indices = {} + + if self.joints_contiguous: + self._frequency_slices["joint"] = slice(selected_joint_ids[0], selected_joint_ids[-1] + 1) + else: + self._frequency_indices["joint"] = wp.array(selected_joint_ids, dtype=int, device=self.device) + + if self.joint_dofs_contiguous: + self._frequency_slices["joint_dof"] = slice(selected_joint_dof_ids[0], selected_joint_dof_ids[-1] + 1) + else: + self._frequency_indices["joint_dof"] = wp.array(selected_joint_dof_ids, dtype=int, device=self.device) + + if self.joint_coords_contiguous: + self._frequency_slices["joint_coord"] = slice(selected_joint_coord_ids[0], selected_joint_coord_ids[-1] + 1) + else: + self._frequency_indices["joint_coord"] = wp.array(selected_joint_coord_ids, dtype=int, device=self.device) + + if self.links_contiguous: + self._frequency_slices["body"] = slice(selected_link_ids[0], selected_link_ids[-1] + 1) + else: + self._frequency_indices["body"] = wp.array(selected_link_ids, dtype=int, device=self.device) + + self.articulation_indices = wp.array(articulation_ids, dtype=int, device=self.device) + + # TODO: zero-stride mask would use less memory + self.full_mask = wp.full(articulation_count, True, dtype=bool, device=self.device) + + # create articulation mask + self.articulation_mask = wp.zeros(model.articulation_count, dtype=bool, device=self.device) + wp.launch( + set_model_articulation_mask_kernel, + dim=articulation_count, + inputs=[self.full_mask, self.articulation_indices, self.articulation_mask], + device=self.device, + ) + + if verbose: + print(f"Articulation '{pattern}': {self.count}") + print(f" Link count: {self.link_count} ({'strided' if self.links_contiguous else 'indexed'})") + print(f" Joint count: {self.joint_count} ({'strided' if self.joints_contiguous else 'indexed'})") + print( + f" DOF count: {self.joint_dof_count} ({'strided' if self.joint_dofs_contiguous else 'indexed'})" + ) + print(f" Fixed base? {self.is_fixed_base}") + print(f" Floating base? {self.is_floating_base}") + print("Link names:") + print(f" {self.body_names}") + print("Joint names:") + print(f" {self.joint_names}") + print("Joint DOF names:") + print(f" {self.joint_dof_names}") + + # ======================================================================================== + # Generic attribute API + + @functools.lru_cache(maxsize=None) # noqa + def _get_attribute_array(self, name: str, source: Model | State | Control, _slice: Slice | None = None): + # get the attribute array + attrib = getattr(source, name) + assert isinstance(attrib, wp.array) + + # reshape with batch dim at front + assert attrib.shape[0] % self.count == 0 + batched_shape = (self.count, attrib.shape[0] // self.count, *attrib.shape[1:]) + attrib = attrib.reshape(batched_shape) + + if _slice is None: + frequency = self.model.get_attribute_frequency(name) + _slice = self._frequency_slices.get(frequency) + else: + _slice = _slice.get() + + if _slice is not None: + # create strided array + attrib = attrib[:, _slice] + else: + # create indexed array + contiguous staging array + _indices = self._frequency_indices.get(frequency) + attrib = wp.indexedarray(attrib, [None, _indices]) + attrib._staging_array = wp.empty_like(attrib) + + return attrib + + def _get_attribute_values(self, name: str, source: Model | State | Control, _slice: slice | None = None): + attrib = self._get_attribute_array(name, source, _slice=_slice) + if hasattr(attrib, "_staging_array"): + wp.copy(attrib._staging_array, attrib) + return attrib._staging_array + else: + return attrib + + # def _set_attribute_values(self, attrib, values, mask=None): + def _set_attribute_values( + self, name: str, target: Model | State | Control, values, mask=None, _slice: slice | None = None + ): + attrib = self._get_attribute_array(name, target, _slice=_slice) + + if not is_array(values) or values.dtype != attrib.dtype: + values = wp.array(values, dtype=attrib.dtype, shape=attrib.shape, device=self.device, copy=False) + assert values.shape == attrib.shape + assert values.dtype == attrib.dtype + + # early out for in-place modifications + if isinstance(attrib, wp.array) and isinstance(values, wp.array): + if values.ptr == attrib.ptr: + return + if isinstance(attrib, wp.indexedarray) and isinstance(values, wp.indexedarray): + if values.data.ptr == attrib.data.ptr: + return + + # get mask + if mask is None: + mask = self.full_mask + else: + if not isinstance(mask, wp.array): + mask = wp.array(mask, dtype=bool, shape=(self.count,), device=self.device, copy=False) + assert mask.shape == (self.count,) + + # launch appropriate kernel based on attribute dimensionality + # TODO: cache concrete overload per attribute? + if attrib.ndim == 1: + wp.launch(set_articulation_attribute_1d_kernel, dim=attrib.shape, inputs=[mask, values, attrib]) + elif attrib.ndim == 2: + wp.launch(set_articulation_attribute_2d_kernel, dim=attrib.shape, inputs=[mask, values, attrib]) + elif attrib.ndim == 3: + wp.launch(set_articulation_attribute_3d_kernel, dim=attrib.shape, inputs=[mask, values, attrib]) + elif attrib.ndim == 4: + wp.launch(set_articulation_attribute_4d_kernel, dim=attrib.shape, inputs=[mask, values, attrib]) + else: + raise NotImplementedError(f"Unsupported attribute with ndim={attrib.ndim}") + + def get_attribute(self, name: str, source: Model | State | Control): + return self._get_attribute_values(name, source) + + def set_attribute(self, name: str, target: Model | State | Control, values, mask=None): + self._set_attribute_values(name, target, values, mask=mask) + + # ======================================================================================== + # Convenience wrappers to align with legacy tensor API + + 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.is_floating_base: + attrib_slice = Slice(self._arti_joint_coord_begin, self._arti_joint_coord_begin + 7) + attrib = self._get_attribute_values("joint_q", source, _slice=attrib_slice) + else: + attrib_slice = Slice(self._arti_joint_begin, self._arti_joint_begin + 1) + attrib = self._get_attribute_values("joint_X_p", self.model, _slice=attrib_slice) + + if attrib.dtype is wp.transform: + return attrib + else: + return wp.array(attrib, dtype=wp.transform, device=self.device, copy=False) + + def set_root_transforms(self, target: Model | State, values: wp.array, mask=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). + values (array): The root transforms to set (dtype=wp.transform). + mask (array): Mask of articulations in this ArticulationView (all by default). + """ + if self.is_floating_base: + attrib_slice = Slice(self._arti_joint_coord_begin, self._arti_joint_coord_begin + 7) + self._set_attribute_values("joint_q", target, values, mask=mask, _slice=attrib_slice) + else: + attrib_slice = Slice(self._arti_joint_begin, self._arti_joint_begin + 1) + self._set_attribute_values("joint_X_p", self.model, values, mask=mask, _slice=attrib_slice) + + 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.is_floating_base: + attrib_slice = Slice(self._arti_joint_dof_begin, self._arti_joint_dof_begin + 6) + attrib = self._get_attribute_values("joint_qd", source, _slice=attrib_slice) + else: + # FIXME? Non-floating articulations have no root velocities. + return None + + if attrib.dtype is wp.spatial_vector: + return attrib + else: + return wp.array(attrib, dtype=wp.spatial_vector, device=self.device, copy=False) + + def set_root_velocities(self, target: Model | State, values: wp.array, mask=None): + """ + Set the root velocities of the articulations. + + Args: + target (Model | State): Where to set the root velocities (Model or State). + values (array): The root velocities to set (dtype=wp.spatial_vector). + mask (array): Mask of articulations in this ArticulationView (all by default). + """ + if self.is_floating_base: + attrib_slice = Slice(self._arti_joint_dof_begin, self._arti_joint_dof_begin + 6) + self._set_attribute_values("joint_qd", target, values, mask=mask, _slice=attrib_slice) + else: + return # no-op + + def get_link_transforms(self, source: Model | State): + return self._get_attribute_values("body_q", source) + + def get_link_velocities(self, source: Model | State): + return self._get_attribute_values("body_qd", source) + + def get_dof_positions(self, source: Model | State): + return self._get_attribute_values("joint_q", source) + + def set_dof_positions(self, target: Model | State, values, mask=None): + self._set_attribute_values("joint_q", target, values, mask=mask) + + def get_dof_velocities(self, source: Model | State): + return self._get_attribute_values("joint_qd", source) + + def set_dof_velocities(self, target: Model | State, values, mask=None): + self._set_attribute_values("joint_qd", target, values, mask=mask) + + def get_dof_forces(self, source: Control): + return self._get_attribute_values("joint_f", source) + + def set_dof_forces(self, target: Control, values, mask=None): + self._set_attribute_values("joint_f", target, values, mask=mask) + + # ======================================================================================== + # Utilities + + def get_model_articulation_mask(self, mask=None): + """ + Get Model articulation mask from a mask in this ArticulationView. + + Args: + mask (array): Mask of articulations in this ArticulationView (all by default). + """ + if mask is None: + return self.articulation_mask + else: + if not isinstance(mask, wp.array): + mask = wp.array(mask, dtype=bool, device=self.device, copy=False) + assert mask.shape == (self.count,) + articulation_mask = wp.zeros(self.model.articulation_count, dtype=bool, device=self.device) + wp.launch( + set_model_articulation_mask_kernel, + dim=mask.size, + inputs=[mask, self.articulation_indices, articulation_mask], + ) + return articulation_mask + + def eval_fk(self, target: Model | State, mask=None): + """ + Evaluates forward kinematics given the joint coordinates and updates the body information. + + Args: + mask (array): Mask of articulations in this ArticulationView (all by default). + """ + # translate view mask to Model articulation mask + articulation_mask = self.get_model_articulation_mask(mask=mask) + newton.sim.eval_fk(self.model, target.joint_q, target.joint_qd, target, mask=articulation_mask)