Selection API

[nvlukasz]

Description

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

[mmacklin]

I think we could move this to the core.articulation module (where we keep other stuff like eval_fk() etc). Would be good to have a collection of articulation functions there.

[nvlukasz]

Makes sense, it should be easy to refactor. I was just trying to keep all the new stuff localized and not touch too many files.

[nvlukasz]

I think it would be useful to have a variant of core.articulation.eval_fk() that takes an array of articulation indices instead of a mask. Then we could get rid of these mask creation kernels. @eric-heiden's GH-148 tweaks the signature of eval_fk() to make the mask optional. Maybe we could add another optional indices arg?

core.articulation.eval_fk(..., mask=mask)  # update from mask
core.articulation.eval_fk(..., indices=indices)  # update from indices

It's a common pattern in Isaac Lab to update a subset of articulations based on the env indices that need to be reset.

[eric-heiden]

That's a good suggestion, I've created an issue for this: #154

[nvlukasz]

Thanks, I could do it once GH-148 is merged.

[Milad-Rakhsha-NV]

yep exactly that is the issue; IL currently uses indices but that isn't graph friendly
What I did in my prototypes was to re-write the IL environments to take a mask array instead to avoid syncing and making everything run in a single graph with physics

[mmacklin]

Yeah, seems like we should try to migrate towards masks in general for this kind of thing since dynamic index sets will always be slow to generate and require host synchronization to count number of instances. That's exactly why I made eval_fk() take a mask originally.

[nvlukasz]

Ok, thanks for all the great feedback. I'll need to change the attribute setter(s) also to take a mask.

It may be best to support both indices and masks for now to give IL some time to transition?

[ghost]

Yes I agree if its possible to support both for the meantime that would be great. We can work towards moving to mask instead of indices meanwhile.

[nvlukasz]

Addressed in #186

[mmacklin]

I think this is look great Lukasz, really simple code, nice clean API, great work!

Before we can merge we probably just need to address a few of these minor comments and add some unit tests.

Thanks.

[Milad-Rakhsha-NV]

It would've been great if we didn't have to put this burden on the user side and could resolve this if/else internally;
Does this also mean that these if/else statements would be propagated all the way to IL to support various solvers

[nvlukasz]

These if-statements are in the examples only to show two different ways of doing things. Normally, as a user, you would pick one and stick to it - no need to provide both code paths and if-statements.

Some background info...

There are two ways of dealing with articulation root transforms:

1. Include the free joint when getting/setting joint_q and joint_qd. This allows you to set the root and dof states in a single call. I think this is the preferred way of doing it in Newton, but it deviates from the legacy tensor API.
2. Set the root states separately from the internal dof states. That requires two calls, but it mimics the legacy tensor API.

I'd be happy to collapse this into a single option, but I think that would mean option 1 (include free joint).

Another thing you might find weird is that the joint_f array used to set joint forces/torques includes the free joint in Newton. So it can be used to apply joint forces and forces on the articulation root in one call, but that again deviates from the legacy tensor API.

So the include_free_joint option is there to ease porting environments from the legacy tensor API to Newton. I'm open to revisiting this. For example, we could always include the free joint. It's the right thing to do for Newton, but it would mean a more complicated rewrite for Isaac Lab.

Or maybe we can keep the include_free_joint flag but make it more granular, so that you can set it on every attribute getter/setter. E.g.,

joint_q = view.get_attribute("joint_q", state, include_free_joint=True)
joint_f = view.get_attribute("joint_f", control, include_free_joint=False)

The flexibility is nice, but it adds complexity.

[nvlukasz]

The more I think about it, the more I want to get rid of this include_free_joint logic. And I would also like to get rid of the special methods for dealing with articulation roots (get/set_root_transforms() and get/set_root_velocities().

I think the selection API should always include the root joint. That is the Newton Way.

This would make it slightly more difficult to port envs from the legacy tensor API, but hear me out. I don't think it'll be that bad.

If the root joint is a free joint:

• To get/set root transforms, use the joint_q attribute slice [:, :7].
• To get/set the internal joint transforms, use the joint_q attribute slice [:, 7:].
• To get/set root velocities, use the joint_qd attribute slice [:, :6].
• To get/set the internal joint velocities, use the joint_qd attribute slice [:, 6:].
• To apply internal joint forces, use the joint_f attribute slice [:, 6:].

So it's essentially up to the user (or Isaac Lab) to read/write the correct slice of the joint arrays. This would simplify the selection API and get rid of the sketchy include_free_joint flag that could be misunderstood or misused. That flag is only there to emulate compatibility with how we did things in PhysX.

@Milad-Rakhsha-NV @mmacklin what do you think?

[nvlukasz]

This should also make it easier to accommodate root joints that are not free joints. We'd just need to tweak the slice offsets based on how many coords/dofs/axes the root joint has.

[nvlukasz]

As an aside, fixed root joints are a bit of a conundrum right now, because they are stripped by the builder before the selection API sees them. So we might need some placeholder joint, something that tells us that there was a fixed joint. Especially if we want to allow getting/setting the transform of the root joint.

Cartpole is a good existing example. According to the joints in the Model, the cart joint is the root joint and the cart is the root link. But that's obviously not correct. We can't change the root transform of the rail, because its shapes are part of the world geometry and there's no easy way to make the descendant links depend on the transform of those static shapes (as far as I can tell).

[ghost]

@nvlukasz how about we have a 2D mask for environments and joints? Anyway we wanted a way to set a subset of joints based on joint indices. If we can have the mask when getting and setting the attributes. Wouldn’t that satisfy both use cases and avoid duplication?

[nvlukasz]

Hmmm, maybe. I don't recall discussing a mask for subsets of joints, just remapping indices based on some orderings.

I think we can do masks, yes, but that doesn't really solve all root transform/velocity issues. For example, if the root joint is a free joint, you can get the root link transform from joint_q. But if it's a fixed joint, you can't get it from joint_q, but have to use joint_X_p. So the masks and attributes you use will need to change depending on the root joint type, which I think will be pretty annoying. That's why I'd rather provide special methods for that, like get_root_transforms() and get_root_velocities().

We can still use masks for joints, but I don't think it fixes everything.

[mmacklin]

I want to clarify a few things about root transforms / free joints / maximal coordinates since there's a lot of issues flying around this atm:

@gyeomannvidia

• Both MuJoCo + Featherstone require a FREE joint for floating base articulations
• XPBD and other maximal coordinate solvers support FREE joints but do not require them to be present, in which case the corresponding root DOFs won't appear in the joint_q array
• Currently when a free joint is present its joint_q DOFs are initialized to zero instead of inheriting the body xform, which can be unintuitive Bug in setting start pose of a rigid body. #93
• body_q is a function of joint_q, and eval_fk() is the way that we sync between them, and must be called before simulation, also related to issue Bug in setting start pose of a rigid body. #93
• Non floating base rigid bodies do not have the root transform stored in the joint_q array (e.g.: revolute base) but are stored in the body_q state and computed using the parent_xform of the joint.

My suggestion is that for Newton we should always add FREE joints for floating base articulations and rigid bodies. This is necessary to allow users to automatically switch between different solvers. (#43). TBD where to do this, finalize() maybe a bit late, so I would suggest importers should handle this for now and that we document clearly this requirement for users building models by hand.

In addition, we should initialize joint_q for FREE joints to the child body xform when calling add_joint_free() so that the previously created body xform is respected.

@nvlukasz

    get_root_transforms() and set_root_transforms()
    get_root_velocities() and set_root_velocities()

Assuming that the ArticulationView knows what the root body is, then I think these root quantities can be retrieved directly from body_q[root_body], body_qd[root_body], in a unified way that doesn't depend on variations in joint layout (e.g.: revolute versus free joint at the root).

[nvlukasz]

Assuming that the ArticulationView knows what the root body is, then I think these root quantities can be retrieved directly from body_q[root_body], body_qd[root_body]

That's what I do currently. I use body_q and body_qd for getting the root link states, but I set them through joint_q and joint_qd if there's a free joint (I had assumed that body_q and body_qd are essentially read-only attributes for articulations).

But I think that using body_q and body_qd is actually wrong. We wanted this API to support reading and writing directly to the Newton buffers without copying. The getters return arrays that alias the Newton buffers with appropriate offsets and strides. Example:

my_joint_q = view.get_attribute("joint_q", state)
wp.launch(magic, ..., outputs=[my_joint_q])

The kernel writes directly to State.joint_q using the strided alias my_joint_q. There are no copies or staging buffers, which is what we wanted. Technically there's no need for a setter, because the data are already in the right place, but there's a plot twist coming and I'll come back to it shortly (*).

Now. What should get_root_transforms() and get_joint_transforms() do? I think they should do something equivalent. Example:

my_roots = view.get_root_transforms(state)
my_joints = view.get_joint_transforms(state)
wp.launch(more_magic, ..., outputs=[my_roots, my_joints])

The kernel writes directly to State.joint_q, this time using two strided aliases, one for the roots and one for the rest of the joints.

• get_root_transforms() == get_attribute("joint_q")[:, :7]
• get_joint_transforms() == get_attribute("joint_q")[:, 7:]

I think this is nice and it fulfills what Isaac Lab is asking for.

If get_root_transforms() returned an alias of the body_q array, I don't think the zero-copy approach will work. I assume that we would still need to copy the root transforms to joint_q, otherwise they'll get overwritten on the next simulation step or eval_fk() and the changes will be lost. Please correct me if I'm wrong.

Things get interesting when the root joint is not a free joint. I would still like to support a zero-copy approach to setting all the transforms. For example, if the root joint is fixed:

• get_root_transforms() == get_attribute("joint_X_p")[:, 0] # root joint frame in parent
• get_joint_transforms() == get_attribute("joint_q") # all internal joint coords

Is this correct? I'm basing it on some sample code from @eric-heiden that used joint_q for free joints and joint_X_p for fixed joints.

I would like to keep this zero-copy mode as much as possible, though it may not be possible in all scenarios.

Zero-copy may not be possible if we want to re-order the joints or use masks like @oahmednv requested. We may need staging buffers to swizzle the joint data around.

my_joint_q = view.get_attribute("joint_q", state)

If the user requested a different joint order when constructing the ArticulationView, then the ordering of joints in my_joint_q is different than in State.joint_q. We have a couple of possibilities here.

1. Create a contiguous staging buffer for this and copy reordered joint data from joint_q into it.
2. Return a wp.indexedarray that remaps the joint axes as needed.

Creating a staging buffer is a con, but the pro is that it's compatible with PyTorch and other array-based frameworks. That's basically what we did in the previous tensor APIs. Creating a wp.indexedarray over State.joint_q avoids a staging buffer. The indexed array can be used in zero-copy mode from Warp, but for PyTorch we still need to create a contiguous copy. And even with Warp, introducing wp.indexedarray into the mix can be a curve ball. If in some cases we return a strided wp.array (no reordering) and in others we return a wp.indexedarray, then the user must be prepared for it. That means writing kernels that support both, which brings new caveats.

(*) Getting back to setters. With zero-copy mode, setters are not needed, because we just write to the Newton arrays directly. But as described above, zero-copy is not always possible, so a setter is needed to put the data in the right place. The API guidance I recommend is that the setter should always be called. Under the hood, we can early-out if the values are already in the right place, so it will be essentially free. But for API correctness, a setter should be called. Furthermore, if we ever need to notify the solver that something has changed, the setter is the place to do it.

So yeah, designing this API is quite challenging with all the different objectives that sometimes contradict each other.

• Support zero-copy mode, especially with Warp, but also with PyTorch et al if possible.
• Support joint reordering or masks.
• Stay true to the Newtonic ways - the "raw" attribute API accomplishes that I think.
• Stay compatible with previous tensor APIs - that's what methods like get_root_transforms() should accomplish.

So we offer two API flavours, one for working with the raw attributes and one that provides an interface that resembles the previous tensor APIs. With raw attributes, you can:

# get root and all joints together
transforms = view.get_attribute("joint_q", state)
velocities = view.get_attribute("joint_qd", state)
...
# set root and all joints together
view.set_attribute("joint_q", state, transforms)
view.set_attribute("joint_qd", state, velocities)

And equivalently, with the other methods:

# get root and internal joints separately
root_transforms = view.get_root_transforms(state)
root_velocities = view.get_root_velocities(state)
joint_transforms = view.get_joint_transforms(state)
joint_velocities = view.get_joint_velocities(state)
...
# set root and internal joints separately
view.set_root_transforms(state, root_transforms)
view.set_root_velocities(state, root_velocities)
view.set_joint_transforms(state, joint_transforms)
view.set_joint_velocities(state, joint_velocities)

With the notable difference that if we change the root free joint to a fixed joint, the code that uses raw attributes will need to change, but the code that uses the other methods can stay the same. So the raw attributes are a more "low-level" way to work with Newton, and the other methods create a light abstraction layer.

[nvlukasz]

Alright, I think I found good solutions. There are some breaking changes, so I created a separate PR for review: #186

[AntoineRichard]

Should be: self.articulation_mask = wp.zeros(model.articulation_count, dtype=bool, device=self.device)

[nvlukasz]

Thanks, nice catch!

[nvlukasz]

Fixed