RobotStatePredictionClient.cpp

Go to the documentation of this file.

/*
 * This file is part of ArmarX.
 *
 * ArmarX is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * ArmarX is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 *
 * @package    RobotAPI::ArmarXObjects::RobotStatePredictionClientExample
 * @author     Rainer Kartmann ( rainer dot kartmann at kit dot edu )
 * @date       2022
 * @copyright  http://www.gnu.org/licenses/gpl-2.0.txt
 *             GNU General Public License
 */
 
#include "RobotStatePredictionClient.h"
 
#include <ArmarXCore/core/exceptions/local/ExpressionException.h>
 
#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
{
 
    RobotStatePredictionClient::RobotStatePredictionClient()
    {
    }
 
 
    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(coreSegmentName))
                .providerSegments(qf::all())
                .entities(qf::all())
                .snapshots(qf::latest());
        }
        armem::client::QueryResult result = remote.reader.query(qb);
        if (result.success)
        {
            return armem::getEntityIDs(result.memory);
        }
        else
        {
            // ToDo: Use exceptions to escalate error.
            ARMARX_VERBOSE << "Query failed: " << result.errorMessage;
            return {};
        }
    }
 
 
    std::vector<armem::PredictionRequest>
    RobotStatePredictionClient::makePredictionRequests(
        const std::vector<armem::MemoryID>& entityIDs,
        armem::Time predictedTime,
        const std::string& engineID)
    {
        std::vector<armem::PredictionRequest> requests;
        requests.reserve(entityIDs.size());
        for (const armem::MemoryID& entityID : entityIDs)
        {
            armem::PredictionRequest& request = requests.emplace_back();
            request.snapshotID = entityID.withTimestamp(predictedTime);
            request.predictionSettings.predictionEngineID = engineID;
        }
        ARMARX_CHECK_EQUAL(requests.size(), entityIDs.size());
        return requests;
    }
 
 
    std::vector<PredictionResult>
    RobotStatePredictionClient::predict(const std::vector<MemoryID>& entityIDs,
                                        Time predictedTime,
                                        const std::string& engineID)
    {
        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;
    }
 
 
    std::optional<Eigen::Affine3f>
    RobotStatePredictionClient::lookupGlobalPose(
        const std::vector<armem::PredictionResult>& localizationPredictionResults,
        const std::vector<armem::MemoryID>& localizationEntityIDs,
        armem::Time predictedTime)
    {
        ARMARX_CHECK_EQUAL(localizationPredictionResults.size(), localizationEntityIDs.size());
        std::stringstream errorMessages;
 
        namespace loc = armem::common::robot_state::localization;
 
        std::optional<armem::wm::CoreSegment> coreSegment;
 
        loc::TransformQuery query;
        query.header.parentFrame = armarx::GlobalFrame;
        query.header.frame = armem::robot_state::constants::robotRootNodeName;
        query.header.agent = "";
        query.header.timestamp = predictedTime;
 
        for (size_t i = 0; i < localizationPredictionResults.size(); ++i)
        {
            const armem::PredictionResult& result = localizationPredictionResults[i];
            const armem::MemoryID& entityID = localizationEntityIDs.at(i);
 
            if (result.success)
            {
                if (query.header.agent.empty())
                {
                    const arondto::Transform tf = arondto::Transform::FromAron(result.prediction);
                    query.header.agent = tf.header.agent;
                }
                if (not coreSegment)
                {
                    coreSegment.emplace(entityID.getCoreSegmentID());
                }
                {
                    armem::EntityUpdate update;
                    update.entityID = entityID;
                    update.referencedTime = predictedTime;
                    update.instancesData = {result.prediction};
                    coreSegment->update(update);
                }
            }
            else
            {
                errorMessages << "Prediction of '" << entityID << "'"
                              << " failed: " << result.errorMessage << "\n";
            }
        }
 
        std::optional<Eigen::Affine3f> result;
        if (coreSegment.has_value())
        {
            loc::TransformResult tf =
                loc::TransformHelper::lookupTransform(coreSegment.value(), query);
 
            if (tf)
            {
                result = tf.transform.transform;
            }
            else
            {
                errorMessages << "\nFailed to lookup transform: " << tf.errorMessage << "\n";
            }
        }
 
        if (not errorMessages.str().empty())
        {
            // ToDo: Introduce an exception here?
            ARMARX_VERBOSE << errorMessages.str();
        }
        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,
                                                  Time predictedTime,
                                                  const std::string& engineID)
    {
        const std::vector<PredictionResult> results = predict(entityIDs, predictedTime, engineID);
        std::optional<Eigen::Affine3f> pose = lookupGlobalPose(results, entityIDs, predictedTime);
        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