Expanded reset logic for Simulator (#2451)

* refactor to add expanded reset logic for Simulator including: reset object states, set light setup default key to scene's light key

* update the internal reset logic to add zero init for positions and velocities when no scene_instance init was provided. Also update unit test.

* refactor C++ tests leveraging reset for time resets only.

src/esp/physics/PhysicsManager.h

@@ -264,11 +264,17 @@ class PhysicsManager : public std::enable_shared_from_this<PhysicsManager> {
 
   /**
    * @brief Reset the simulation and physical world.
-   * Sets the @ref worldTime_ to 0.0, does not change physical state.
+   * Sets the @ref worldTime_ to 0.0, changes the physical state of all objects back to their initial states. Only changes motion_type when scene_instance specified a motion type.
    */
   virtual void reset() {
-    /* TODO: reset object states or clear them? Other? */
+    // reset object states from initial values (e.g. from scene instance)
     worldTime_ = 0.0;
+    for (const auto& bro : existingObjects_) {
+      bro.second->resetStateFromSceneInstanceAttr();
+    }
+    for (const auto& bao : existingArticulatedObjects_) {
+      bao.second->resetStateFromSceneInstanceAttr();
+    }
   }
 
   /** @brief Stores references to a set of drawable elements. */

src/esp/physics/RigidObject.h

@@ -207,8 +207,7 @@ class RigidObject : public RigidBase {
   VelocityControl::ptr getVelocityControl() { return velControl_; };
 
   /**
-   * @brief Set the object's state from a @ref
-   * esp::metadata::attributes::SceneObjectInstanceAttributes
+   * @brief Set the object's state from a @ref esp::metadata::attributes::SceneObjectInstanceAttributes
    */
   void resetStateFromSceneInstanceAttr() override;
 

src/esp/physics/bullet/BulletArticulatedObject.cpp

@@ -279,8 +279,23 @@ void BulletArticulatedObject::resetStateFromSceneInstanceAttr() {
   if (attrObjMotionType != physics::MotionType::UNDEFINED) {
     setMotionType(attrObjMotionType);
   }
-  // set initial joint positions
-  // get array of existing joint dofs
+
+  // initially set positions to zero (or identity quaternion for spherical
+  // joints)
+  int posCount = 0;
+  // init unit quat
+  float quat_init[] = {0, 0, 0, 1};
+  for (int i = 0; i < btMultiBody_->getNumLinks(); ++i) {
+    auto& link = btMultiBody_->getLink(i);
+    if (link.m_posVarCount > 0) {
+      // feed the unit quat to all setters, only spherical joints will hit the
+      // one at the end
+      btMultiBody_->setJointPosMultiDof(i, const_cast<float*>(quat_init));
+      posCount += link.m_posVarCount;
+    }
+  }
+
+  // set initial joint positions from instance config if applicable
   std::vector<float> aoJointPose = getJointPositions();
   // get instance-specified initial joint positions
   const auto& initJointPos = sceneObjInstanceAttr->getInitJointPose();
@@ -295,8 +310,9 @@ void BulletArticulatedObject::resetStateFromSceneInstanceAttr() {
   }
   setJointPositions(aoJointPose);
 
+  // first clear all joint vels
+  setJointVelocities(std::vector<float>(size_t(btMultiBody_->getNumDofs())));
   // set initial joint velocities
-  // get array of existing joint vel dofs
   std::vector<float> aoJointVels = getJointVelocities();
   // get instance-specified initial joint velocities
   std::vector<float> initJointVels =
@@ -311,6 +327,10 @@ void BulletArticulatedObject::resetStateFromSceneInstanceAttr() {
     }
     aoJointVels[i] = initJointVels[i];
   }
+  setJointVelocities(aoJointVels);
+
+  // clear any forces
+  setJointForces(std::vector<float>(size_t(btMultiBody_->getNumDofs())));
 
 }  // BulletArticulatedObject::resetStateFromSceneInstanceAttr
 

src/esp/sim/Simulator.cpp

@@ -672,7 +672,9 @@ bool Simulator::instanceArticulatedObjectsForSceneAttributes(
 
 void Simulator::reset() {
   if (physicsManager_ != nullptr) {
-    // Note: only resets time to 0 by default.
+    // Note: resets time to 0 and all existing objects set back to initial
+    // states. Does not add back deleted objects or delete added objects. Does
+    // not break ManagedObject pointers.
     physicsManager_->reset();
   }
 

src/esp/sim/Simulator.h

@@ -68,6 +68,14 @@ class Simulator {
 
   void reconfigure(const SimulatorConfiguration& cfg);
 
+  /**
+   * @brief Reset the simulation state including the state of all physics
+   * objects, agents, and the default light setup.
+   * Sets the @ref worldTime_ to 0.0, changes the physical state of all objects back to their initial states.
+   * Does not invalidate existing ManagedObject wrappers.
+   * Does not add or remove object instances.
+   * Only changes motion_type when scene_instance specified a motion type.
+   */
   void reset();
 
   void seed(uint32_t newSeed);

src/tests/PhysicsTest.cpp

@@ -648,13 +648,12 @@ void PhysicsTest::testRemoveSleepingSupport() {
     cubes[0]->setMotionType(esp::physics::MotionType::STATIC);
 
     for (int testCase = 0; testCase < 2; ++testCase) {
-      // reset time to 0, should not otherwise modify state
-      physicsManager_->reset();
+      float currentTime = physicsManager_->getWorldTime();
       CORRADE_COMPARE_AS(physicsManager_->getNumRigidObjects(), 0,
                          Cr::TestSuite::Compare::Greater);
 
       // simulate to stabilize the stack and populate collision islands
-      while (physicsManager_->getWorldTime() < 4.0) {
+      while (physicsManager_->getWorldTime() < currentTime + 4.0) {
         physicsManager_->stepPhysics(0.1);
       }
 
@@ -924,15 +923,15 @@ void PhysicsTest::testVelocityControl() {
     objectWrapper->resetTransformation();
     objectWrapper->setTranslation(Magnum::Vector3{0, 2.0, 0});
     physicsManager_->setGravity({});  // 0 gravity interference
-    physicsManager_->reset();         // reset time to 0
 
     // should closely follow kinematic result while uninhibited in 0 gravity
     float targetTime = 0.5;
+    float currentTime = physicsManager_->getWorldTime();
     esp::core::RigidState initialObjectState(objectWrapper->getRotation(),
                                              objectWrapper->getTranslation());
     esp::core::RigidState kinematicResult =
         velControl->integrateTransform(targetTime, initialObjectState);
-    while (physicsManager_->getWorldTime() < targetTime) {
+    while (physicsManager_->getWorldTime() < currentTime + targetTime) {
       physicsManager_->stepPhysics(physicsManager_->getTimestep());
     }
     CORRADE_COMPARE_AS(
@@ -946,7 +945,7 @@ void PhysicsTest::testVelocityControl() {
 
     // should then get blocked by ground plane collision
     targetTime = 2.0;
-    while (physicsManager_->getWorldTime() < targetTime) {
+    while (physicsManager_->getWorldTime() < currentTime + targetTime) {
       physicsManager_->stepPhysics(physicsManager_->getTimestep());
     }
     CORRADE_COMPARE_AS(objectWrapper->getTranslation()[1], 1.0 - errorEps,
@@ -964,8 +963,8 @@ void PhysicsTest::testVelocityControl() {
   velControl->linVelIsLocal = true;
 
   targetTime = 10.0;
-  physicsManager_->reset();  // reset time to 0
-  while (physicsManager_->getWorldTime() < targetTime) {
+  float currentTime = physicsManager_->getWorldTime();
+  while (physicsManager_->getWorldTime() < currentTime + targetTime) {
     physicsManager_->stepPhysics(physicsManager_->getTimestep());
   }
 

src_python/habitat_sim/simulator.py

@@ -172,6 +172,17 @@ def reset(self, agent_ids: Optional[int] = None) -> ObservationDict:
     def reset(
         self, agent_ids: Union[Optional[int], List[int]] = None
     ) -> Union[ObservationDict, Dict[int, ObservationDict],]:
+        """
+        Reset the simulation state including the state of all physics objects, agents, and the default light setup.
+        Sets the world time to 0.0, changes the physical state of all objects back to their initial states.
+        Does not invalidate existing ManagedObject wrappers.
+        Does not add or remove object instances.
+        Only changes motion_type when scene_instance specified a motion type.
+
+        :param agent_ids: An optional list of agent ids for which to return the sensor observations. If none is provide, default agent is used.
+
+        :return: Sensor observations in the reset state.
+        """
         super().reset()
         for i in range(len(self.agents)):
             self.reset_agent(i)

tests/test_physics.py

@@ -447,7 +447,7 @@ def test_velocity_control():
         obj_handle = obj_template_mgr.get_template_handle_by_id(template_ids[0])
 
         for iteration in range(2):
-            sim.reset()
+            cur_time = sim.get_world_time()
 
             box_object = rigid_obj_mgr.add_object_by_template_handle(obj_handle)
             # Verify is not articulated
@@ -471,17 +471,19 @@ def test_velocity_control():
             vel_control.controlling_lin_vel = True
             vel_control.controlling_ang_vel = True
 
-            while sim.get_world_time() < 1.0:
+            while sim.get_world_time() < 1.0 + cur_time:
                 # NOTE: stepping close to default timestep to get near-constant velocity control of DYNAMIC bodies.
                 sim.step_physics(0.00416)
 
-            ground_truth_pos = sim.get_world_time() * vel_control.linear_velocity
+            ground_truth_pos = (
+                sim.get_world_time() - cur_time
+            ) * vel_control.linear_velocity
             assert np.allclose(box_object.translation, ground_truth_pos, atol=0.01)
             ground_truth_q = mn.Quaternion([[0, 0.480551, 0], 0.876967])
             angle_error = mn.math.half_angle(ground_truth_q, box_object.rotation)
             assert angle_error < mn.Rad(0.005)
 
-            sim.reset()
+            cur_time = sim.get_world_time()
 
             # test local velocities (turn in a half circle)
             vel_control.lin_vel_is_local = True
@@ -492,15 +494,12 @@ def test_velocity_control():
             box_object.translation = [0.0, 0.0, 0.0]
             box_object.rotation = mn.Quaternion()
 
-            while sim.get_world_time() < 0.5:
+            while sim.get_world_time() < 0.5 + cur_time:
                 # NOTE: stepping close to default timestep to get near-constant velocity control of DYNAMIC bodies.
                 sim.step_physics(0.008)
 
-            print(sim.get_world_time())
-
             # NOTE: explicit integration, so expect some error
             ground_truth_q = mn.Quaternion([[1.0, 0.0, 0.0], 0.0])
-            print(box_object.translation)
             assert np.allclose(
                 box_object.translation, np.array([0, 1.0, 0.0]), atol=0.07
             )

tests/test_simulator.py

@@ -74,24 +74,92 @@ def test_sim_reset(make_cfg_settings):
     mm = habitat_sim.metadata.MetadataMediator(hab_cfg.sim_cfg)
     hab_cfg_mm.metadata_mediator = mm
 
+    def check_isclose(val1, val2):
+        return np.isclose(val1, val2, rtol=1e-4).all()
+
     test_list = [hab_cfg, hab_cfg_mm]
     for ctor_arg in test_list:
         with habitat_sim.Simulator(ctor_arg) as sim:
             agent_config = sim.config.agents[0]
             sim.initialize_agent(0)
+            # cache agent initial state
             initial_state = sim.agents[0].initial_state
             # Take random steps in the environment
             for _ in range(10):
                 action = random.choice(list(agent_config.action_space.keys()))
                 sim.step(action)
 
+            # add rigid and articulated objects
+            sim.metadata_mediator.ao_template_manager.load_configs(
+                "data/test_assets/urdf/"
+            )
+            sim.metadata_mediator.object_template_manager.load_configs(
+                "data/test_assets/objects/"
+            )
+
+            chair_handle = (
+                sim.metadata_mediator.object_template_manager.get_template_handles(
+                    "chair"
+                )[0]
+            )
+            ao_handle = sim.metadata_mediator.ao_template_manager.get_template_handles(
+                "prism"
+            )[0]
+
+            ro = sim.get_rigid_object_manager().add_object_by_template_handle(
+                chair_handle
+            )
+            ao = sim.get_articulated_object_manager().add_articulated_object_by_template_handle(
+                ao_handle
+            )
+
+            assert ro is not None
+            assert ao is not None
+
+            # cache the initial state for verification
+            ao_initial_state = (
+                ao.transformation,
+                ao.joint_positions,
+                ao.joint_velocities,
+            )
+            ro_initial_state = ro.transformation
+
+            assert check_isclose(ro.transformation, ro_initial_state)
+            assert check_isclose(ao.transformation, ao_initial_state[0])
+            assert check_isclose(ao.joint_positions, ao_initial_state[1])
+            assert check_isclose(ao.joint_velocities, ao_initial_state[2])
+
+            ro.translation = mn.Vector3(1, 2, 3)
+            ro.rotation = mn.Quaternion.rotation(
+                mn.Rad(0.123), mn.Vector3(0.1, 0.2, 0.3).normalized()
+            )
+            ao.translation = mn.Vector3(3, 2, 1)
+            ao.rotation = mn.Quaternion.rotation(
+                mn.Rad(0.321), mn.Vector3(0.3, 0.2, 0.1).normalized()
+            )
+            ao.joint_positions = np.array(ao_initial_state[1]) * 0.2
+            ao.joint_velocities = np.ones(len(ao.joint_velocities))
+
+            assert not check_isclose(ro.transformation, ro_initial_state)
+            assert not check_isclose(ao.transformation, ao_initial_state[0])
+            assert not check_isclose(ao.joint_positions, ao_initial_state[1])
+            assert not check_isclose(ao.joint_velocities, ao_initial_state[2])
+
+            # do the reset
             sim.reset()
+
+            # validate agent state reset
             new_state = sim.agents[0].get_state()
-            same_position = all(initial_state.position == new_state.position)
-            same_rotation = np.isclose(
-                initial_state.rotation, new_state.rotation, rtol=1e-4
-            )  # Numerical error can cause slight deviations
-            assert same_position and same_rotation
+            # NOTE: Numerical error can cause slight deviations, use isclose
+            # assert same agent position and rotation
+            assert check_isclose(initial_state.position, new_state.position)
+            assert check_isclose(initial_state.rotation, new_state.rotation)
+
+            # validate object state resets
+            assert check_isclose(ro.transformation, ro_initial_state)
+            assert check_isclose(ao.transformation, ao_initial_state[0])
+            assert check_isclose(ao.joint_positions, ao_initial_state[1])
+            assert check_isclose(ao.joint_velocities, ao_initial_state[2])
 
 
 def test_sim_multiagent_move_and_reset(make_cfg_settings, num_agents=10):