[Simulation Interfaces] Support for reset state (#875)

* Support for reset state

* Fix issue with mixing local and global transform

---------

Signed-off-by: Michał Pełka <[email protected]>

Gems/SimulationInterfaces/Code/Include/SimulationInterfaces/SimulationEntityManagerRequestBus.h

@@ -96,6 +96,10 @@ namespace SimulationInterfaces
             const AZ::Transform& initialPose,
             const bool allowRename,
             SpawnCompletedCb completedCb) = 0;
+
+        //! Reset the simulation to begin.
+        //! This will revert the entire simulation to the initial state.
+        virtual void ResetAllEntitiesToInitialState() = 0;
     };
 
     class SimulationInterfacesBusTraits : public AZ::EBusTraits

Gems/SimulationInterfaces/Code/Source/Clients/SimulationEntitiesManager.cpp

@@ -157,6 +157,16 @@ namespace SimulationInterfaces
                     spawnData->second.m_completedCb(AZ::Success(registeredName));
                     m_spawnCompletedCallbacks.erase(spawnData);
                 }
+
+                // cache the initial state
+                EntityState initialState{};
+                initialState.m_pose = entity->GetTransform()->GetLocalTM();
+                if (rigidBody)
+                {
+                    initialState.m_twist_linear = rigidBody->GetLinearVelocity();
+                    initialState.m_twist_angular = rigidBody->GetAngularVelocity();
+                }
+                m_entityIdToInitialState[entityId] = initialState;
             });
         m_simulationBodyRemovedHandler = AzPhysics::SceneEvents::OnSimulationBodyRemoved::Handler(
             [this](AzPhysics::SceneHandle sceneHandle, AzPhysics::SimulatedBodyHandle bodyHandle)
@@ -273,6 +283,11 @@ namespace SimulationInterfaces
             m_entityIdToSimulatedEntityMap.erase(findIt);
             m_simulatedEntityToEntityIdMap.erase(simulatedEntityName);
         }
+        auto findIt2 = m_entityIdToInitialState.find(entityId);
+        if (findIt2 != m_entityIdToInitialState.end())
+        {
+            m_entityIdToInitialState.erase(findIt2);
+        }
     }
 
     AZ::Outcome<EntityNameList, FailedResult> SimulationEntitiesManager::GetEntities(const EntityFilters& filter)
@@ -417,17 +432,23 @@ namespace SimulationInterfaces
         // get entity and all descendants
         AZStd::vector<AZ::EntityId> entityAndDescendants;
         AZ::TransformBus::EventResult(entityAndDescendants, entityId, &AZ::TransformBus::Events::GetEntityAndAllDescendants);
+        SetEntitiesState(entityAndDescendants, state);
+        return AZ::Success();
+    }
 
+    void SimulationEntitiesManager::SetEntitiesState(const AZStd::vector<AZ::EntityId>& entityAndDescendants, const EntityState& state)
+    {
+        AZ_Assert(!entityAndDescendants.empty(), "Entity and descendants list is empty");
+        if (entityAndDescendants.empty())
+        {
+            return;
+        }
+        const AZ::EntityId entityId = entityAndDescendants.front();
         if (state.m_pose.IsOrthogonal())
         {
             // disable simulation for all entities
-            AZStd::map<AZ::EntityId, AZ::Transform> entityTransforms;
             for (const auto& descendant : entityAndDescendants)
             {
-                // get name
-                AZStd::string entityName = "Unknown";
-                AZ::ComponentApplicationBus::BroadcastResult(entityName, &AZ::ComponentApplicationRequests::GetEntityName, descendant);
-                AZ_Printf("SimulationInterfaces", "Disable physics for entity %s\n", entityName.c_str());
                 Physics::RigidBodyRequestBus::Event(descendant, &Physics::RigidBodyRequests::DisablePhysics);
             }
 
@@ -451,7 +472,6 @@ namespace SimulationInterfaces
                 SetRigidBodyVelocities(rigidBody, state);
             }
         }
-        return AZ::Success();
     }
 
     void SimulationEntitiesManager::DeleteEntity(const AZStd::string& name, DeletionCompletedCb completedCb)
@@ -480,9 +500,7 @@ namespace SimulationInterfaces
         const auto ticketId = entity->GetEntitySpawnTicketId();
         if (m_spawnedTickets.find(ticketId) != m_spawnedTickets.end())
         {
-            // remove the ticket
-            // m_spawnedTickets.erase(ticketId);
-            /// get spawner
+            // get spawner
             auto spawner = AZ::Interface<AzFramework::SpawnableEntitiesDefinition>::Get();
             AZ_Assert(spawner, "SpawnableEntitiesDefinition is not available.");
             // get ticket
@@ -637,7 +655,6 @@ namespace SimulationInterfaces
             {
                 return;
             }
-            const AZ::Entity* root = *view.begin();
 
             for (auto* entity : view)
             {
@@ -663,11 +680,19 @@ namespace SimulationInterfaces
                 entity->SetName(entityName);
             }
 
-            auto* transformInterface = root->FindComponent<AzFramework::TransformComponent>();
-            if (transformInterface)
+            // get the first entity
+            if (view.size()>1)
             {
-                transformInterface->SetWorldTM(initialPose);
+                const AZ::Entity* firstEntity = view[1];
+                AZ_Assert(firstEntity, "First entity is not available");
+                auto* transformInterface = firstEntity->FindComponent<AzFramework::TransformComponent>();
+                if (transformInterface)
+                {
+                    transformInterface->SetWorldTM(initialPose);
+                }
             }
+
+
         };
         optionalArgs.m_completionCallback =
             [this, uri](AzFramework::EntitySpawnTicket::Id ticketId, AzFramework::SpawnableConstEntityContainerView view)
@@ -725,4 +750,15 @@ namespace SimulationInterfaces
         }
         return simulatedEntityName;
     }
+
+    void SimulationEntitiesManager::ResetAllEntitiesToInitialState()
+    {
+        for (const auto& [entityId, initialState] : m_entityIdToInitialState)
+        {
+            AZStd::vector<AZ::EntityId> entityAndDescendants;
+            AZ::TransformBus::EventResult(entityAndDescendants, entityId, &AZ::TransformBus::Events::GetEntityAndAllDescendants);
+
+            SetEntitiesState(entityAndDescendants, initialState);
+        }
+    }
 } // namespace SimulationInterfaces

Gems/SimulationInterfaces/Code/Source/Clients/SimulationEntitiesManager.h

@@ -51,6 +51,7 @@ namespace SimulationInterfaces
             const AZ::Transform& initialPose,
             const bool allowRename,
             SpawnCompletedCb completedCb) override;
+        void ResetAllEntitiesToInitialState() override;
 
         // AZ::Component interface implementation
         void Init() override;
@@ -58,6 +59,8 @@ namespace SimulationInterfaces
         void Deactivate() override;
 
     private:
+
+
         //! Registers simulated entity to entity id mapping.
         //! Note that the entityId will be registered under unique name.
         //! \param entityId The entity id to register
@@ -71,6 +74,9 @@ namespace SimulationInterfaces
         //! Returns the simulated entity name for the given entity id.
         AZStd::string GetSimulatedEntityName(AZ::EntityId entityId, const AZStd::string& proposedName) const;
 
+        //! Set the state of the entity and their descendants.
+        void SetEntitiesState(const AZStd::vector<AZ::EntityId>& entityAndDescendants, const EntityState& state);
+
         AzPhysics::SceneEvents::OnSimulationBodyAdded::Handler m_simulationBodyAddedHandler;
         AzPhysics::SceneEvents::OnSimulationBodyRemoved::Handler m_simulationBodyRemovedHandler;
 
@@ -79,8 +85,7 @@ namespace SimulationInterfaces
         AzPhysics::SceneHandle m_physicsScenesHandle = AzPhysics::InvalidSceneHandle;
         AZStd::unordered_map<AZStd::string, AZ::EntityId> m_simulatedEntityToEntityIdMap;
         AZStd::unordered_map<AZ::EntityId, AZStd::string> m_entityIdToSimulatedEntityMap;
-        AZStd::unordered_set<AzPhysics::SimulatedBodyHandle> m_disabledBodies;
-
+        AZStd::unordered_map<AZ::EntityId, EntityState> m_entityIdToInitialState;
         AZStd::unordered_map<AzFramework::EntitySpawnTicket::Id, AzFramework::EntitySpawnTicket> m_spawnedTickets;
 
         struct SpawnCompletedCbData

Gems/SimulationInterfaces/Code/Source/Clients/SimulationManager.cpp

@@ -93,7 +93,7 @@ namespace SimulationInterfaces
             &SimulationFeaturesAggregatorRequests::AddSimulationFeatures,
             AZStd::unordered_set<SimulationFeatures>{ simulation_interfaces::msg::SimulatorFeatures::SIMULATION_RESET,
                                                       simulation_interfaces::msg::SimulatorFeatures::SIMULATION_RESET_TIME,
-                                                      // simulation_interfaces::msg::SimulatorFeatures::SIMULATION_RESET_STATE,
+                                                      simulation_interfaces::msg::SimulatorFeatures::SIMULATION_RESET_STATE,
                                                       simulation_interfaces::msg::SimulatorFeatures::SIMULATION_RESET_SPAWNED,
                                                       simulation_interfaces::msg::SimulatorFeatures::SIMULATION_STATE_PAUSE,
                                                       simulation_interfaces::msg::SimulatorFeatures::STEP_SIMULATION_SINGLE,

Gems/SimulationInterfacesROS2/Code/Source/Services/ResetSimulationServiceHandler.cpp

@@ -23,7 +23,7 @@ namespace SimulationInterfacesROS2
     {
         return AZStd::unordered_set<AZ::u8>{ SimulationFeatures::SIMULATION_RESET,
                                              SimulationFeatures::SIMULATION_RESET_TIME,
-                                             //SimulationFeatures::SIMULATION_RESET_STATE,
+                                             SimulationFeatures::SIMULATION_RESET_STATE,
                                              SimulationFeatures::SIMULATION_RESET_SPAWNED};
     }
 
@@ -33,8 +33,9 @@ namespace SimulationInterfacesROS2
         if (request.scope == simulation_interfaces::srv::ResetSimulation::Request::SCOPE_STATE)
         {
             Response response;
-            response.result.result = simulation_interfaces::msg::Result::RESULT_FEATURE_UNSUPPORTED;
-            response.result.error_message = "Not implemented yet";
+            SimulationInterfaces::SimulationEntityManagerRequestBus::Broadcast(
+                &SimulationInterfaces::SimulationEntityManagerRequests::ResetAllEntitiesToInitialState);
+            response.result.result = simulation_interfaces::msg::Result::RESULT_OK;
             return response;
         }
 

Gems/SimulationInterfacesROS2/Code/Tests/Tools/Mocks/SimulationEntityManagerMock.h

@@ -24,6 +24,7 @@ namespace UnitTest
         MOCK_METHOD2(DeleteEntity, void(const AZStd::string& name, DeletionCompletedCb completedCb));
         MOCK_METHOD1(DeleteAllEntities, void(DeletionCompletedCb completedCb));
         MOCK_METHOD0(GetSpawnables, AZ::Outcome<SpawnableList, FailedResult>());
+        MOCK_METHOD0(ResetAllEntitiesToInitialState, void());
         MOCK_METHOD6(
             SpawnEntity,
             void(const AZStd::string& name, const AZStd::string& uri, const AZStd::string& entityNamespace, const AZ::Transform& initialPose, const bool allowRename, SpawnCompletedCb completedCb));