parallel update for get_obstacles_from_stage().get_collision_check_world()

[zoctipus]

So far, for a multi-envs simulation, I think one of the largest slow down is non-parallel(sequential) update of get_collision_check_world().
Since collision world is actually List[WorldConfig] in batched environment where each index in the list corresponds to each WorldConfig of each env. I wonder if there's any parallel implementation for updating all indices of this List[WorldConfig]. I am currently using forloop call of get_obstacles_from_stage(only_path = env[idx])

something I tried but didn't work is
helper function

def _get_env_collsion_config(idx) -> WorldConfig:
     return self.usd_help.get_obstacles_from_stage(
                ignore_substring = ['Robot', 'panda'],
                only_paths = [f"/World/envs/env_{}".format(idx)],
                reference_prim_path = f'/World/envs/env_{}/robot_{}/panda_link0'.format(idx)
            ).get_collision_check_world()

Parallel implementation attempt:

futures:List[torch.Future] = []
for idx in range(self.num_robots):
    furture_job = torch.fork(_get_env_collsion_config, idx)
    futures.append(furture_job)
world_cfg_list = []
for future in futures:
    world_cfg_list.append(torch.wait(future))

the implementation fails because at line
furture_job = torch.fork(_get_env_collsion_config, idx)
_get_env_collsion return type is WorldConfig, not a valid type for Future Class(Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type WorldConfig.)

I wonder if theres any parallel implementation for update WorldConfig_list, if theres not, could you plan in future release? I'd also love to help writing such code if help is needed(though I need some advice)!
Thank you for you hard work.

[balakumar-s]

The slowdown is because of cuRobo's current use of USD stage parsing to read simulation state. This can be slow when querying for many environments. For every environment, we access the robot in the USD stage and query poses of objects w.r.t. robot.

We want to move to USDRT (https://docs.omniverse.nvidia.com/kit/docs/usdrt/latest/docs/usd_fabric_usdrt.html) which can be significantly faster for parsing simulation state. However, we haven't yet investigated it thoroughly.

Ideally, we would create a class similar to USDHelper in https://github.com/NVlabs/curobo/blob/main/src/curobo/util/usd_helper.py that can parse a USD stage much faster than the current implementation.