Spawning and Managing Multiple Articulations

[btx0424]

Hi team, you have been doing absolutely wonderful work in developing this framework. This post is to raise some questions and thoughts regarding spawning and managing multiple articulations in an environment (with the devel branch).

Expected practice to spawn an articulation multiple times for each environment

Although it is possible to just use different Articulation instances, merging homogeneous ones could improve the performance. Since it is currently impossible to do this by specifying prim paths like "/World/envs/env_*/Robot-[1,2]", I wonder what the expected best practice is now and in the future.

For example:

Option 1: allow users to pass a sequence of prim paths

pros: straightforward, concise in code
cons: more complicated logic for view creation

Option 2: create articulations separately and create a view afterward

pros: does not need further modification
cons: need a flag to avoid repeated view creation

Managing multiple articulations

Currently, the articulations are indexed only by env_ids, which effectively assumes there is only one instance of this articulation in each environment. This poses some inconvenience for implementing tasks with multiple homogeneous agents. A possible way to expand this behavior is to introduce shapes to the views. For example, an articulation view corresponds to "/World/envs/env_*/Robot-[1,2]" has a shape of [num_envs, 2]. Then the APIs look like:

articulation: Articulation = ...
articulation[env_ids, entity_ids].reset()
articulation[env_ids, entity_ids].write_*_to_sim(...)

obs = articulation[:, :2].data.root_pos_w
# or
obs = articulation.data[:, :2].root_pos_w

where the indices are resolved internally.

This of course involves a lot of modifications. I am humbly posting this both as a user and a prospective contributor. WHYT?

[Mayankm96]

Hi @btx0424 ,

Thank you for bringing up this discussion!

We considered both options (when we were rewriting the core pieces) and decided to go with (2) since it simplified the asset spawning and data indexing.

While I understand the motivation of (1), the main question is how much performance gain we will get with that. Making everything "batched" within the environment may require resolving a lot of indexing, which I feel could be counter-productive.

What has your experience been?

[btx0424]

Hi @Mayankm96 ,

Thanks for replying. I can see your point since advanced indexing in PyTorch usually makes tensor copies instead of views.

However, I think one important advantage of "batching" articulation management is that it saves the MDP logic from for loops. Otherwise, we would have to write something like for i in range(self.num_agents) in every action, observation, and reward function. If there are, say, 2 to 3 agents, this is fine. But the overhead grows as the agent number increases and there is also overhead of concatenating things together. I think it is also useful for a manipulation task with many objects to compute observation/reward with.

Regarding index resolving, it seems if the articulations are homogeneous, no extra overhead is introduced except for _reset_idx which is relatively sparse. Specifically, we could have:

class Articulation(...):
  def __init__(self, ..., shape):
    ...
    self._shape = self._ALL_INDICES.reshape(shape).shape
  
  def __getitem__(self, index = slice(None)):
    # do some checking
    index = self._ALL_INDICES.reshape(self._shape)[index]
    # read and write using `index`
    return _ArticulationView(self, index)

A regular step operates on all instances, i.e., env_ids==slice(None), so it's effectively self._ALL_INDICES.reshape(self._shape)[...]. Everything is still contiguous and no copies are created.