Add object "getter" utils to sim_utilities (facebookresearch#1816)

* utils for getting all scene object instances

* add ao_link_map getter 

* add single object getters from id and handle

* add tests for all getter features

habitat-lab/habitat/sims/habitat_simulator/sim_utilities.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import magnum as mn
 
@@ -302,7 +302,12 @@ def snap_down(
 def get_all_object_ids(sim: habitat_sim.Simulator) -> Dict[int, str]:
     """
     Generate a dict mapping all active object ids to a descriptive string containing the object instance handle and, for ArticulatedLinks, the link name.
+
+    :param sim: The Simulator instance.
+
+    :return: a dict mapping object ids to a descriptive string.
     """
+
     rom = sim.get_rigid_object_manager()
     aom = sim.get_articulated_object_manager()
 
@@ -321,3 +326,111 @@ def get_all_object_ids(sim: habitat_sim.Simulator) -> Dict[int, str]:
             )
 
     return object_id_map
+
+
+def get_all_objects(
+    sim: habitat_sim.Simulator,
+) -> List[
+    Union[
+        habitat_sim.physics.ManagedRigidObject,
+        habitat_sim.physics.ManagedArticulatedObject,
+    ]
+]:
+    """
+    Get a list of all ManagedRigidObjects and ManagedArticulatedObjects in the scene.
+
+    :param sim: The Simulator instance.
+
+    :return: a list of ManagedObject wrapper instances containing all objects currently instantiated in the scene.
+    """
+
+    managers = [
+        sim.get_rigid_object_manager(),
+        sim.get_articulated_object_manager(),
+    ]
+    all_objects = []
+    for mngr in managers:
+        all_objects.extend(mngr.get_objects_by_handle_substring().values())
+    return all_objects
+
+
+def get_ao_link_id_map(sim: habitat_sim.Simulator) -> Dict[int, int]:
+    """
+    Construct a dict mapping ArticulatedLink object_id to parent ArticulatedObject object_id.
+    NOTE: also maps ao's root object id to itself for ease of use.
+
+    :param sim: The Simulator instance.
+
+    :return: dict mapping ArticulatedLink object ids to parent object ids.
+    """
+
+    aom = sim.get_articulated_object_manager()
+    ao_link_map: Dict[int, int] = {}
+    for ao in aom.get_objects_by_handle_substring().values():
+        # add the ao itself for ease of use
+        ao_link_map[ao.object_id] = ao.object_id
+        # add the links
+        for link_id in ao.link_object_ids:
+            ao_link_map[link_id] = ao.object_id
+
+    return ao_link_map
+
+
+def get_obj_from_id(
+    sim: habitat_sim.Simulator,
+    obj_id: int,
+    ao_link_map: Optional[Dict[int, int]] = None,
+) -> Union[
+    habitat_sim.physics.ManagedRigidObject,
+    habitat_sim.physics.ManagedArticulatedObject,
+]:
+    """
+    Get a ManagedRigidObject or ManagedArticulatedObject from an object_id.
+
+    ArticulatedLink object_ids will return the ManagedArticulatedObject.
+    If you want link id, use ManagedArticulatedObject.link_object_ids[obj_id].
+
+    :param sim: The Simulator instance.
+    :param obj_id: object id for which ManagedObject is desired.
+    :param ao_link_map: A pre-computed map from link object ids to their parent ArticulatedObject's object id.
+
+    :return: a ManagedObject or None
+    """
+
+    if ao_link_map is None:
+        # Note: better to pre-compute this and pass it around
+        ao_link_map = get_ao_link_id_map(sim)
+
+    rom = sim.get_rigid_object_manager()
+    if rom.get_library_has_id(obj_id):
+        return rom.get_object_by_id(obj_id)
+    aom = sim.get_articulated_object_manager()
+    if obj_id in ao_link_map:
+        return aom.get_object_by_id(ao_link_map[obj_id])
+
+    return None
+
+
+def get_obj_from_handle(
+    sim: habitat_sim.Simulator, obj_handle: str
+) -> Union[
+    habitat_sim.physics.ManagedRigidObject,
+    habitat_sim.physics.ManagedArticulatedObject,
+]:
+    """
+    Get a ManagedRigidObject or ManagedArticulatedObject from its instance handle.
+
+    :param sim: The Simulator instance.
+    :param obj_handle: object istance handle for which ManagedObject is desired.
+
+    :return: a ManagedObject or None
+    """
+
+    rom = sim.get_rigid_object_manager()
+    if rom.get_library_has_handle(obj_handle):
+        return rom.get_object_by_handle(obj_handle)
+    aom = sim.get_articulated_object_manager()
+    if aom.get_library_has_handle(obj_handle):
+        return aom.get_object_by_handle(obj_handle)
+
+    return None

test/test_sim_utils.py

@@ -11,6 +11,11 @@
 
 from habitat.sims.habitat_simulator.sim_utilities import (
     bb_ray_prescreen,
+    get_all_object_ids,
+    get_all_objects,
+    get_ao_link_id_map,
+    get_obj_from_handle,
+    get_obj_from_id,
     snap_down,
 )
 from habitat_sim import Simulator, built_with_bullet
@@ -155,3 +160,65 @@ def test_snap_down(support_margin, obj_margin, stage_support):
                 if stage_support:
                     # don't need 3 iterations for stage b/c no motion types to test
                     break
+
+
+@pytest.mark.skipif(
+    not built_with_bullet,
+    reason="Raycasting API requires Bullet physics.",
+)
+@pytest.mark.skipif(
+    not osp.exists("data/replica_cad/"),
+    reason="Requires ReplicaCAD dataset.",
+)
+def test_object_getters():
+    sim_settings = default_sim_settings.copy()
+    sim_settings[
+        "scene_dataset_config_file"
+    ] = "data/replica_cad/replicaCAD.scene_dataset_config.json"
+    sim_settings["scene"] = "apt_0"
+    hab_cfg = make_cfg(sim_settings)
+    with Simulator(hab_cfg) as sim:
+        # scrape various lists from utils
+        all_objects = get_all_objects(sim)
+        all_object_ids = get_all_object_ids(sim)
+        ao_link_map = get_ao_link_id_map(sim)
+
+        # validate parity between util results
+        assert len(all_objects) == (
+            sim.get_rigid_object_manager().get_num_objects()
+            + sim.get_articulated_object_manager().get_num_objects()
+        )
+        for object_id in ao_link_map:
+            assert (
+                object_id in all_object_ids
+            ), f"Link or AO object id {object_id} is not found in the global object id map."
+        for obj in all_objects:
+            assert obj is not None
+            assert obj.is_alive
+            assert (
+                obj.object_id in all_object_ids
+            ), f"Object's object_id {object_id} is not found in the global object id map."
+            # check the wrapper getter functions
+            obj_from_id_getter = get_obj_from_id(
+                sim, obj.object_id, ao_link_map
+            )
+            obj_from_handle_getter = get_obj_from_handle(sim, obj.handle)
+            assert obj_from_id_getter.object_id == obj.object_id
+            assert obj_from_handle_getter.object_id == obj.object_id
+
+        # specifically validate link object_id mapping
+        aom = sim.get_articulated_object_manager()
+        for ao in aom.get_objects_by_handle_substring().values():
+            assert ao.object_id in all_object_ids
+            assert ao.object_id in ao_link_map
+            link_indices = ao.get_link_ids()
+            assert len(ao.link_object_ids) == len(link_indices)
+            for link_object_id, link_index in ao.link_object_ids.items():
+                assert link_object_id in ao_link_map
+                assert ao_link_map[link_object_id] == ao.object_id
+                assert link_index in link_indices
+                # links should return reference to parent object
+                obj_from_id_getter = get_obj_from_id(
+                    sim, link_object_id, ao_link_map
+                )
+                assert obj_from_id_getter.object_id == ao.object_id