Add predictions of joint positions

source/RobotAPI/components/armem/client/RobotStatePredictionClientExample/CMakeLists.txt

@@ -24,6 +24,8 @@ set(SOURCES
 set(HEADERS
     # ToDo: Move to library.
     RobotStatePredictionClient.h
+    # ToDo: Move to simox.
+    simox_alg.hpp
 
     Component.h
     Impl.h

source/RobotAPI/components/armem/client/RobotStatePredictionClientExample/Impl.cpp

@@ -31,12 +31,83 @@
 #include <RobotAPI/libraries/armem/client/query.h>
 #include <RobotAPI/libraries/armem/core/error.h>
 #include <RobotAPI/libraries/armem/core/wm/memory_definitions.h>
+#include <RobotAPI/libraries/armem_robot_state/aron/Proprioception.aron.generated.h>
 #include <RobotAPI/libraries/armem_robot_state/aron/Transform.aron.generated.h>
 #include <RobotAPI/libraries/armem_robot_state/client/common/VirtualRobotReader.h>
 #include <RobotAPI/libraries/armem_robot_state/client/common/constants.h>
 #include <RobotAPI/libraries/armem_robot_state/common/localization/TransformHelper.h>
 
 
+namespace simox::alg
+{
+    template <class... Args>
+    std::vector<Args...>
+    concatenate(const std::vector<Args...>& lhs, const std::vector<Args...>& rhs)
+    {
+        std::vector<Args...> conc = lhs;
+        std::copy(rhs.begin(), rhs.end(), std::back_inserter(conc));
+        return conc;
+    }
+
+
+    template <class KeyT, class ValueT>
+    std::map<KeyT, ValueT>
+    map_from_key_value_pairs(const std::vector<KeyT>& lhs, const std::vector<ValueT>& rhs)
+    {
+        const size_t size = std::min(lhs.size(), rhs.size());
+
+        std::map<KeyT, ValueT> map;
+        for (size_t i = 0; i < size; ++i)
+        {
+            map.emplace(lhs[i], rhs[i]);
+        }
+        return map;
+    }
+
+
+    template <class KeyT, class ValueT>
+    std::vector<ValueT>
+    multi_at(const std::map<KeyT, ValueT>& map,
+             const std::vector<KeyT>& keys,
+             bool skipMissing = false)
+    {
+        std::vector<ValueT> values;
+        values.reserve(keys.size());
+
+        for (const KeyT& key : keys)
+        {
+            if (skipMissing)
+            {
+                if (auto it = map.find(key); it != map.end())
+                {
+                    values.push_back(it->second);
+                }
+            }
+            else
+            {
+                // Throw an exception if missing.
+                values.push_back(map.at(key));
+            }
+        }
+
+        return values;
+    }
+
+    template <class... Args>
+    std::vector<Args...>
+    slice(const std::vector<Args...>& vector,
+          size_t start = 0,
+          std::optional<size_t> end = std::nullopt)
+    {
+        std::vector<Args...> result;
+        auto beginIt = vector.begin() + start;
+        auto endIt = end ? vector.begin() + *end : vector.end();
+        std::copy(beginIt, endIt, std::back_inserter(result));
+        return result;
+    }
+
+} // namespace simox::alg
+
 namespace armarx::robot_state_prediction_client_example
 {
 
@@ -141,37 +212,49 @@ namespace armarx::robot_state_prediction_client_example
             return;
         }
 
-        const std::vector<armem::MemoryID> locEntityIDs = localizationEntityIDs.has_value()
-                                                              ? localizationEntityIDs.value()
+        // Which entities to predict?
+
+        const std::vector<armem::MemoryID> locEntityIDs = this->localizationEntityIDs.has_value()
+                                                              ? this->localizationEntityIDs.value()
                                                               : client.queryLocalizationEntityIDs();
-        if (not locEntityIDs.empty())
-        {
-            // Predict.
+        const std::vector<armem::MemoryID> propEntityIDs =
+            this->propioceptionEntityIDs.has_value() ? this->propioceptionEntityIDs.value()
+                                                     : client.queryProprioceptionEntityIDs();
 
-            const std::optional<Eigen::Affine3f> globalPose =
-                client.predictGlobalPose(locEntityIDs, predictedTime, "Linear");
+        // Predict.
 
-            if (globalPose.has_value())
-            {
-                // ARMARX_INFO << "Predicted global pose: \n" << globalPose->matrix();
-                robotViz->pose(globalPose->matrix());
+        auto prediction = client.predictWholeBody(
+            locEntityIDs, propEntityIDs, predictedTime, properties.robotName, "Linear");
 
-                if (not localizationEntityIDs)
-                {
-                    // Store entity IDs for successful lookup.
-                    localizationEntityIDs = locEntityIDs;
-                }
-            }
+        // Gather results.
 
+        if (prediction.globalPose.has_value())
+        {
+            // ARMARX_INFO << "Predicted global pose: \n" << globalPose->matrix();
+            robotViz->pose(prediction.globalPose->matrix());
 
-            // Visualize.
+            if (not localizationEntityIDs)
+            {
+                // Store entity IDs for successful lookup.
+                this->localizationEntityIDs = locEntityIDs;
+            }
+        }
+        if (prediction.jointPositions.has_value())
+        {
+            robotViz->joints(prediction.jointPositions.value());
 
+            if (not propioceptionEntityIDs)
             {
-                viz::Layer layer = remote.arviz.layer(properties.robotName + " Prediction");
-                layer.add(robotViz.value());
-                remote.arviz.commit(layer);
+                this->propioceptionEntityIDs = propEntityIDs;
             }
         }
+
+        // Visualize.
+        {
+            viz::Layer layer = remote.arviz.layer(properties.robotName + " Prediction");
+            layer.add(robotViz.value());
+            remote.arviz.commit(layer);
+        }
     }
 
 

source/RobotAPI/components/armem/client/RobotStatePredictionClientExample/Impl.h

@@ -80,6 +80,7 @@ namespace armarx::robot_state_prediction_client_example
         armem::robot_state::RobotStatePredictionClient client;
 
         std::optional<std::vector<armem::MemoryID>> localizationEntityIDs;
+        std::optional<std::vector<armem::MemoryID>> propioceptionEntityIDs;
         std::optional<viz::Robot> robotViz;
 
     };

source/RobotAPI/components/armem/client/RobotStatePredictionClientExample/RobotStatePredictionClient.cpp

@@ -27,13 +27,17 @@
 #include <RobotAPI/libraries/armem/client/MemoryNameSystem.h>
 #include <RobotAPI/libraries/armem/client/query.h>
 #include <RobotAPI/libraries/armem/core/error.h>
+#include <RobotAPI/libraries/armem/core/operations.h>
 #include <RobotAPI/libraries/armem/core/wm/memory_definitions.h>
+#include <RobotAPI/libraries/armem_robot_state/aron/Proprioception.aron.generated.h>
 #include <RobotAPI/libraries/armem_robot_state/aron/Transform.aron.generated.h>
 #include <RobotAPI/libraries/armem_robot_state/aron_conversions.h>
 #include <RobotAPI/libraries/armem_robot_state/client/common/constants.h>
 #include <RobotAPI/libraries/armem_robot_state/common/localization/TransformHelper.h>
 #include <RobotAPI/libraries/core/FramedPose.h>
 
+#include "simox_alg.hpp"
+
 
 namespace armarx::armem::robot_state
 {
@@ -45,11 +49,25 @@ namespace armarx::armem::robot_state
 
     std::vector<armem::MemoryID>
     RobotStatePredictionClient::queryLocalizationEntityIDs()
+    {
+        return _queryEntityIDs(armem::robot_state::constants::localizationCoreSegment);
+    }
+
+
+    std::vector<MemoryID>
+    RobotStatePredictionClient::queryProprioceptionEntityIDs()
+    {
+        return _queryEntityIDs(armem::robot_state::constants::proprioceptionCoreSegment);
+    }
+
+
+    std::vector<MemoryID>
+    RobotStatePredictionClient::_queryEntityIDs(const std::string& coreSegmentName)
     {
         armem::client::QueryBuilder qb;
         {
             namespace qf = armarx::armem::client::query_fns;
-            qb.coreSegments(qf::withName(armem::robot_state::constants::localizationCoreSegment))
+            qb.coreSegments(qf::withName(coreSegmentName))
                 .providerSegments(qf::all())
                 .entities(qf::all())
                 .snapshots(qf::latest());
@@ -57,16 +75,12 @@ namespace armarx::armem::robot_state
         armem::client::QueryResult result = remote.reader.query(qb);
         if (result.success)
         {
-            std::vector<armem::MemoryID> entityIDs;
-
-            result.memory.forEachEntity([&entityIDs](const armem::wm::Entity& entity)
-                                        { entityIDs.push_back(entity.id()); });
-
-            return entityIDs;
+            return armem::getEntityIDs(result.memory);
         }
         else
         {
-            ARMARX_INFO << "Query failed: " << result.errorMessage;
+            // ToDo: Use exceptions to escalate error.
+            ARMARX_VERBOSE << "Query failed: " << result.errorMessage;
             return {};
         }
     }
@@ -98,6 +112,13 @@ namespace armarx::armem::robot_state
     {
         const std::vector<PredictionRequest> requests =
             makePredictionRequests(entityIDs, predictedTime, engineID);
+        return predict(requests);
+    }
+
+
+    std::vector<PredictionResult>
+    RobotStatePredictionClient::predict(const std::vector<PredictionRequest>& requests)
+    {
         const std::vector<PredictionResult> results = remote.reader.predict(requests);
         ARMARX_CHECK_EQUAL(results.size(), requests.size());
         return results;
@@ -178,6 +199,27 @@ namespace armarx::armem::robot_state
         return result;
     }
 
+    std::optional<std::map<std::string, float>>
+    RobotStatePredictionClient::lookupJointPositions(
+        const std::vector<PredictionResult>& proprioceptionPredictionResults,
+        const std::string& robotName)
+    {
+        auto it = std::find_if(proprioceptionPredictionResults.begin(),
+                               proprioceptionPredictionResults.end(),
+                               [&robotName](const PredictionResult& result)
+                               { return result.snapshotID.entityName == robotName; });
+        if (it != proprioceptionPredictionResults.end())
+        {
+            auto result = *it;
+            auto prop = armem::arondto::Proprioception::FromAron(result.prediction);
+            return prop.joints.position;
+        }
+        else
+        {
+            return std::nullopt;
+        }
+    }
+
 
     std::optional<Eigen::Affine3f>
     RobotStatePredictionClient::predictGlobalPose(const std::vector<MemoryID>& entityIDs,
@@ -189,4 +231,50 @@ namespace armarx::armem::robot_state
         return pose;
     }
 
+
+    std::optional<std::map<std::string, float>>
+    RobotStatePredictionClient::predictJointPositions(const std::vector<MemoryID>& entityIDs,
+                                                      Time predictedTime,
+                                                      const std::string& robotName,
+                                                      const std::string& engineID)
+    {
+        const std::vector<PredictionResult> results = predict(entityIDs, predictedTime, engineID);
+        return lookupJointPositions(results, robotName);
+    }
+
+
+    RobotStatePredictionClient::WholeBodyPrediction
+    RobotStatePredictionClient::predictWholeBody(const std::vector<armem::MemoryID>& locEntityIDs,
+                                                 const std::vector<armem::MemoryID>& propEntityIDs,
+                                                 Time predictedTime,
+                                                 const std::string& robotName,
+                                                 const std::string& engineID)
+    {
+        RobotStatePredictionClient::WholeBodyPrediction result;
+
+        // Predict.
+        const std::vector<armem::MemoryID> entityIDs =
+            simox::alg::concatenate(locEntityIDs, propEntityIDs);
+
+        if (entityIDs.empty())
+        {
+            return result;
+        }
+
+        auto _results = predict(entityIDs, predictedTime, "Linear");
+
+        // Gather results.
+        std::vector<armem::PredictionResult> locResults, propResults;
+        locResults = simox::alg::slice(_results, 0, locEntityIDs.size());
+        propResults = simox::alg::slice(_results, locEntityIDs.size());
+        ARMARX_CHECK_EQUAL(locResults.size(), locEntityIDs.size());
+        ARMARX_CHECK_EQUAL(propResults.size(), propEntityIDs.size());
+
+
+        result.globalPose = lookupGlobalPose(locResults, locEntityIDs, predictedTime);
+        result.jointPositions = lookupJointPositions(propResults, robotName);
+
+        return result;
+    }
+
 } // namespace armarx::armem::robot_state

source/RobotAPI/components/armem/client/RobotStatePredictionClientExample/RobotStatePredictionClient.h

@@ -47,13 +47,35 @@ namespace armarx::armem::robot_state
                           armem::Time predictedTime,
                           const std::string& engineID = "Linear");
 
+        std::optional<std::map<std::string, float>>
+        predictJointPositions(const std::vector<armem::MemoryID>& entityIDs,
+                              armem::Time predictedTime,
+                              const std::string& robotName,
+                              const std::string& engineID = "Linear");
+
+
+        struct WholeBodyPrediction
+        {
+            std::optional<Eigen::Affine3f> globalPose;
+            std::optional<std::map<std::string, float>> jointPositions;
+        };
+        WholeBodyPrediction
+        predictWholeBody(const std::vector<armem::MemoryID>& localizationEntityIDs,
+                         const std::vector<armem::MemoryID>& proprioceptionEntityIDs,
+                         armem::Time predictedTime,
+                         const std::string& robotName,
+                         const std::string& engineID = "Linear");
+
 
         std::vector<armem::PredictionResult> predict(const std::vector<armem::MemoryID>& entityIDs,
                                                      armem::Time predictedTime,
                                                      const std::string& engineID = "Linear");
+        std::vector<armem::PredictionResult>
+        predict(const std::vector<armem::PredictionRequest>& requests);
 
 
         std::vector<armem::MemoryID> queryLocalizationEntityIDs();
+        std::vector<armem::MemoryID> queryProprioceptionEntityIDs();
 
         std::vector<armem::PredictionRequest>
         makePredictionRequests(const std::vector<armem::MemoryID>& entityIDs,
@@ -65,6 +87,13 @@ namespace armarx::armem::robot_state
                          const std::vector<armem::MemoryID>& localizationEntityIDs,
                          armem::Time predictedTime);
 
+        std::optional<std::map<std::string, float>> lookupJointPositions(
+            const std::vector<armem::PredictionResult>& proprioceptionPredictionResults,
+            const std::string& robotName);
+
+
+    private:
+        std::vector<armem::MemoryID> _queryEntityIDs(const std::string& coreSegmentName);
 
     public:
         struct Remote