Automatic projector for OMPL compound state spaces composed of projectable subspaces

CompoundProjectionEvaluator.h

#ifndef _COMPOUNDPROJECTIONEVALUATOR_H_
#define _COMPOUNDPROJECTIONEVALUATOR_H_

#include <unordered_set>


namespace ob = ompl::base;


std::unordered_set<int> pickSet(int N, int k, std::mt19937& gen)
{
    std::unordered_set<int> elems;
    for (int r = N - k; r < N; ++r) {
        int v = std::uniform_int_distribution<>(1, r)(gen);

        // there are two cases.
        // v is not in candidates ==> add it
        // v is in candidates ==> well, r is definitely not, because
        // this is the first iteration in the loop that we could've
        // picked something that big.

        if (!elems.insert(v).second) {
            elems.insert(r);
        }   
    }
    return elems;
}


std::vector<int> pick(int N, int k) {
    std::random_device rd;
    std::mt19937 gen(rd());

    std::unordered_set<int> elems = pickSet(N, k, gen);

    // ok, now we have a set of k elements. but now
    // it's in a [unknown] deterministic order.
    // so we have to shuffle it:

    std::vector<int> result(elems.begin(), elems.end());
    std::shuffle(result.begin(), result.end(), gen);
    return result;
}


class CompoundProjectionEvaluator: public ob::ProjectionEvaluator
    {
    public:
        CompoundProjectionEvaluator(const ob::StateSpacePtr space,
                                    const int nsample=3);

        void setup() override;

        unsigned int getDimension() const override;

        void project(const ob::State *state, Eigen::Ref<Eigen::VectorXd> projection) const override;

    protected:
        std::vector<ob::ProjectionEvaluatorPtr> proj_;
        unsigned long int ndims_;
        unsigned long int nsubspaces_;
        std::vector<double> joined_cell_sizes;
        unsigned long int nsample_;
        std::vector<int> indices;

    };


CompoundProjectionEvaluator::CompoundProjectionEvaluator(const ob::StateSpacePtr space,
    const int nsample)
  : ProjectionEvaluator(space)
{
    this -> nsample_ = nsample;

    // Ensure the state is compound
    if (!space_->isCompound())
        throw ompl::Exception("Cannot construct an auto subspace projection evaluator for a space that is not compound");
    auto comp = space_->as<ob::CompoundStateSpace>();
    nsubspaces_ = comp->getSubspaceCount(); 

    // Get the default projections for each subspace
    ndims_ = 0;
    for (long unsigned int i=0; i < nsubspaces_; i++){
        ndims_ += comp->getSubspace(i)->getDefaultProjection()->getDimension();
    }

    // Construct random indices into the array
    this->indices = pick(ndims_, nsample_);

    // For each subspace get the default projection and full cell sizes
    std::vector<double> jcs; 
    for (unsigned long int i=0; i<this->nsubspaces_; i++){
        auto sub_proj = space_->as<ob::CompoundStateSpace>()->getSubspace(i)->getDefaultProjection();
        proj_.push_back(sub_proj);
        auto cells = sub_proj->getCellSizes();
        jcs.insert(jcs.end(), 
                    cells.begin(), 
                    cells.end());
    }
    this->joined_cell_sizes = jcs;

    // Sample the cell sizes
    std::vector<double> sampled_cell_sizes;
    for (unsigned long int i=0; i<nsample_; i++){
        auto this_cell = joined_cell_sizes[this->indices[i]];
        sampled_cell_sizes.push_back(this_cell);
    }
    this->setCellSizes(sampled_cell_sizes);
}


void CompoundProjectionEvaluator::setup()
{
    ProjectionEvaluator::setup();
}


unsigned int CompoundProjectionEvaluator::getDimension() const
{
    return this -> nsample_;
}


void CompoundProjectionEvaluator::project(const ob::State *state, Eigen::Ref<Eigen::VectorXd> projection) const
{
    // Set the state as a compound state
    auto comp = state->as<ob::CompoundState>();

    // Unsampled output
    Eigen::VectorXd output(this->ndims_);

    // Fill the output with the subspace projections
    int n = 0;
    for (long unsigned int i=0; i <this->nsubspaces_; i++){

        auto subsize = proj_[i]->getDimension();
        Eigen::VectorXd tmp(subsize);
        proj_[i]->project(comp->components[i], tmp);
        for (long unsigned int j=0; j<subsize; j++){
            output[n] = tmp[j];
            n++;
        }
    }

    // Use the random indices
    for (long unsigned i=0; i<this->nsample_; i++){
        projection[i] = output[this->indices[i]];
    }
}


#endif