allow for parallel rendering.

[branyang02]

Originally, obs.reset() returns:

# policy.arm_joint_pos: Tensor, shape=(7,), dtype=torch.float32, device=cuda:0
# policy.gripper_pos: Tensor, shape=(1,), dtype=torch.float32, device=cuda:0
# policy.external_cam: Tensor, shape=(1, 720, 1280, 3), dtype=torch.uint8, device=cuda:0
# policy.wrist_cam: Tensor, shape=(1, 720, 1280, 3), dtype=torch.uint8, device=cuda:0

Made changes such that obs.reset() returns

# policy.arm_joint_pos: Tensor, shape=(B, 7), dtype=torch.float32, device=cuda:0
# policy.gripper_pos: Tensor, shape=(B, 1), dtype=torch.float32, device=cuda:0
# policy.external_cam: Tensor, shape=(B, 720, 1280, 3), dtype=torch.uint8, device=cuda:0
# policy.wrist_cam: Tensor, shape=(B, 720, 1280, 3), dtype=torch.uint8, device=cuda:0

This is only a fix on the IsaacLab side, and I did not test whether this works or not with the current infer functions. If we do decide to merge this implementation. Please make sure to test the infer functions in jax models!!!

[arhanjain]

thanks for the contrib! currently the openpi policy serving doesn't allow for batched inference. i'll check w openpi folks if its a PR they're willing to take. if so, can merge this once openpi is updated