Add a cache to ObjectStateMachine.

examples/hitl/rearrange_v2/object_state_manipulator.py

@@ -99,7 +99,8 @@ def set_object_state(
         Set an object state, regardless of constraints.
         """
         obj = sim_utilities.get_obj_from_handle(self._sim, object_handle)
-        set_state_of_obj(obj, state_name, state_value)
+        osm = self._sim.object_state_machine
+        set_state_of_obj(obj, state_name, state_value, osm)
 
     def get_boolean_action(
         self, state_name: str, target_value: bool

examples/hitl/rearrange_v2/world.py

@@ -65,8 +65,6 @@ def __init__(
         self._regions: Dict[int, SemanticRegion] = {}
         # Cache of all active states.
         self._all_states: Optional[Dict[str, ObjectStateSpec]] = None
-        # Cache of object states.
-        self._state_snapshot_dict: Dict[str, Dict[str, Any]] = {}
         # Object state container.
         self._object_state_machine: Optional[ObjectStateMachine] = None
         # Cache of categories for each object handles.
@@ -144,27 +142,18 @@ def update(self, dt: float) -> None:
 
         if osm is not None:
             osm.update_states(sim, dt)
-            self._state_snapshot_dict = osm.get_snapshot_dict(sim)
 
     def get_state_snapshot_dict(self) -> Dict[str, Dict[str, Any]]:
         """
         Get a snapshot of the current state of all objects.
-
-        Example:
-        {
-            "is_powered_on": {
-                "my_lamp.0001": True,
-                "my_oven": False,
-                ...
-            },
-            "is_clean": {
-                "my_dish.0002:" False,
-                ...
-            },
-            ...
-        }
         """
-        return self._state_snapshot_dict
+        sim = self._sim
+        osm = self._object_state_machine
+
+        if osm is not None:
+            return osm.get_snapshot_dict(sim)
+        else:
+            return {}
 
     def get_all_states(self) -> Dict[str, ObjectStateSpec]:
         """

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

@@ -6,8 +6,10 @@
 
 """This module implements a singleton state-machine architecture for representing and managing non-geometric object states via metadata manipulation. For example, tracking and manipulating state such as "powered on" or "clean vs dirty". This interface is intended to provide a foundation which can be extended for downstream applications."""
 
+from __future__ import annotations
+
 from collections import defaultdict
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import magnum as mn
 
@@ -47,20 +49,25 @@ def set_state_of_obj(
     obj: Union[ManagedArticulatedObject, ManagedRigidObject],
     state_name: str,
     state_val: Any,
+    object_state_machine: Optional[ObjectStateMachine] = None,
 ) -> None:
     """
     Set the specified state in an object's "object_states" user_defined metadata.
 
     :param obj: The ManagedObject.
     :param state_name: The name/key of the object state property to set.
     :param state_val: The value of the object state property to set.
+    :param object_state_machine: If specified, flag the snapshot as dirty.
     """
 
     user_attr = obj.user_attributes
     obj_state_config = user_attr.get_subconfig("object_states")
     obj_state_config.set(state_name, state_val)
     user_attr.save_subconfig("object_states", obj_state_config)
 
+    if object_state_machine is not None:
+        object_state_machine.set_snapshot_dirty()
+
 
 ##################################################
 # Object state machine implementation
@@ -204,21 +211,6 @@ def draw_state(
 
         draw_object_highlight(obj, debug_line_render, camera_transform, color)
 
-    def toggle(
-        self, obj: Union[ManagedArticulatedObject, ManagedRigidObject]
-    ) -> bool:
-        """
-        Toggles a boolean state, returning the newly set value.
-
-        :param obj: The ManagedObject instance.
-        :return: The new value of the state.
-        """
-
-        cur_state = get_state_of_obj(obj, self.name)
-        new_state = not cur_state
-        set_state_of_obj(obj, self.name, new_state)
-        return new_state
-
 
 class ObjectIsClean(BooleanObjectState):
     """
@@ -269,6 +261,10 @@ def __init__(self, active_states: List[ObjectStateSpec] = None) -> None:
         self.objects_with_states: Dict[
             str, List[ObjectStateSpec]
         ] = defaultdict(lambda: [])
+        # Whether the snapshot cache is dirty.
+        self._dirty = False
+        # Snapshot cache.
+        self._snapshot_cache: Dict[str, Dict[str, Any]] = {}
 
     def initialize_object_state_map(self, sim: habitat_sim.Simulator) -> None:
         """
@@ -299,6 +295,8 @@ def register_object(
                     f"registered state {state} for object {obj.handle}"
                 )
 
+        self.set_snapshot_dirty()
+
     def update_states(self, sim: habitat_sim.Simulator, dt: float) -> None:
         """
         Update all tracked object states for a simulation step.
@@ -317,6 +315,14 @@ def update_states(self, sim: habitat_sim.Simulator, dt: float) -> None:
                 for state in states:
                     state.update_state(sim, obj, dt)
 
+        self.set_snapshot_dirty()
+
+    def set_snapshot_dirty(self) -> None:
+        """
+        Flag the snapshot dict for recalculation the next time `get_snapshot_dict()` is called.
+        """
+        self._dirty = True
+
     def get_snapshot_dict(
         self, sim: habitat_sim.Simulator
     ) -> Dict[str, Dict[str, Any]]:
@@ -340,6 +346,9 @@ def get_snapshot_dict(
             >>>     ...
             >>> }
         """
+        if not self._dirty:
+            return self._snapshot_cache
+
         snapshot: Dict[str, Dict[str, Any]] = defaultdict(lambda: {})
         for object_handle, states in self.objects_with_states.items():
             obj = sutils.get_obj_from_handle(sim, object_handle)
@@ -350,4 +359,7 @@ def get_snapshot_dict(
                     if obj_state is not None
                     else state.default_value()
                 )
-        return dict(snapshot)
+
+        self._snapshot_cache = dict(snapshot)
+        self._dirty = False
+        return self._snapshot_cache