Selection API v2

[nvlukasz]

Description

This PR supersedes the first PR (#119).

Attribute API

The attribute API is similar to the previous version, e.g.:

joint_transforms = view.get_attribute("joint_q", state)
joint_velocities = view.get_attribute("joint_qd", state)
...
view.set_attribute("joint_q", state, joint_transforms)
view.set_attribute("joint_qd", state, joint_velocities)

This is a flexible approach that makes it easy to incorporate custom attributes by name without the need to add new API entry points.

By default, the attribute API always includes all joints, links, and DOFs. This allows users familiar with the Newton model to work efficiently with the underlying data.

Joint and Link Filtering

We now allow filtering joints and links to give users more control over articulation views, e.g.:

ArticulationView(..., include_joints=["hip_*", "ankle_*"], exclude_links=["*left*"])

With this foundation in place, it should be easy to add support for different joint/link ordering like BFS or DFS (in a subsequent PR).

Helper API

For spiritual compatibility with the legacy tensor API and to simplify working with different root joint types, we offer a few convenience methods like this:

view.get_root_transforms(state)
view.get_root_velocities(state)

For floating-base articulations, the root transforms are part of "joint_q". For fixed-base articulations, the root transforms are set through "joint_X_p". The helper methods take care of this so that users don't need to maintain different code paths for different root joint types.

The legacy tensor API didn't have the concept of root joints. Articulations had the same number of DOFs regardless of whether they had a fixed or floating base. We can mimic that behaviour by filtering out free joints like this:

ArticulationView(..., exclude_joint_types=[newton.JOINT_FREE])

This will exclude the free joint data from all attributes, making the number of DOFs the same regardless of whether the articulation has a fixed or floating base. Note that set_root_transforms() and set_root_velocities() will still use "joint_q" and "joint_qd" under the hood for floating-base articulations, but get_attribute("joint_q") and get_attribute("joint_qd") will NOT include the root free joint since it's filtered out.

# filter out free joints
view = ArticulationView(model, "ant", exclude_joint_types=[newton.JOINT_FREE])

view.get_root_transforms(state)  # gets root transforms from joint_q[:, :7]
view.get_dof_positions(state)    # gets DOF positions from joint_q[:, 7:]

view.get_root_velocities(state)  # gets root velocities from joint_qd[:, :6]
view.get_dof_velocities(state)   # gets DOF velocities from joint_qd[:, 6:]

# The leading free DOFs are not included with any attributes (filtered out)
view.get_attribute("joint_limit_lower", model)
view.get_attribute("joint_armature", model)
...

# same with DOF forces etc.
dof_forces = 5.0 - 10.0 * torch.rand((self.num_envs, self.ants.joint_dof_count))
self.ants.set_dof_forces(control, dof_forces)

Zero Copy Semantics

It's possible to work with this API in zero-copy mode:

# get strided alias to the state attribute
joint_positions = view.get_attribute("joint_q", state)
# launch a kernel that modifies the state in-place
wp.launch(update_transforms, ..., outputs=[joint_positions])

It's even possible to do this zero-copy modification in PyTorch. The attributes returned by the API are always strided arrays (not indexed arrays), so PyTorch read and write the data directly.

It's technically not necessary to call a setter, because the data are already written to the right place. But for the sake of consistency, it is required to call the setter:

view.set_attribute("joint_q", state, joint_transforms)

There are two reasons for this:

1. In some cases, modifying an attribute may require notifying the solver. The setter will send this notification if needed.
2. Joint and link filtering (or reordering) may present the joints, DOFs, or links in non-contiguous order. This can be solved using indexed arrays in Warp, but tensor-based frameworks like PyTorch need staging buffers where a swizzled copy of the data is held. In that case, we don't have zero-copy access, so a setter is required to copy the data to the right place.

So as a general rule, calling a setter is required if an attribute is modified. The are built-in optimizations that will avoid copying data if the values were written to the attribute in-place. So zero-copy is achievable even if we call setters.

Examples

There are a few examples that show what the APIs look like in action (newton/examples/example_selection_*).

Newton Migration Guide

Please ensure the migration guide for warp.sim users is up-to-date with the changes made in this MR.

• The migration guide in docs/migration.rst is up-to date

Before your PR is "Ready for review"

• All commits are signed-off to indicate that your contribution adheres to the Developer Certificate of Origin requirements
• I understand that GitHub does not perform any GPU testing of this pull request
• Necessary tests have been added
• Documentation is up-to-date
• Code passes formatting and linting checks with pre-commit run -a