Extending python bindings (externally)

[DavidPL1]

Hi,

I am looking for help with the python bindings.

For some context:

I am currently developing a simulation wrapped around MuJoCo that supports ROS.

My simulation is very similar to Simulate, but the key difference is, that it defines an additional abstraction of a MuJoCo environment (among other things, consisting of an mjModel and mjData object) which is separated from a rendered GUI. This environment abstraction defines a multithreaded lifecycle that includes loading, stepping and communicating with the ROS interfaces.

Now I am looking into adding python bindings with pybind11 to facilitate usage through python in conjunction with Machine Learning. And since you already have very well developed python bindings for MuJoCo classes, enums and functions, it would make sense to just interface them in my bindings. Preferably I'd want to import the official bindings in my code instead of building a forked version.

As far as I can tell, currently the only way to create mjModel and mjData instances in python is by loading them through the bindings, whereas I would need to create them from raw pointers.

For clarity, (re)loading with in my bindings to looks like this:

bool loadModel(const std::string &filename) {
    mju::strcpy_arr(filename_queued_, filename.c_str());
    settings_.load_request = 1; // atomic_int that signals lifecycle should run load logic
    while (settings_.load_request > 0) {
        std::this_thread::sleep_for(std::chrono::milliseconds(5)); // wait for load to complete
    }

    // now the members model_ and data_ hold new mjModel and mjData pointers

     return load_error_[0] != '\0';
} 

And then I'd like to somehow grant access to model_ and data_, where somewhere along the way they get casted to their equivalents in the official python bindings.

Is there any way to do this? I should also mention that I am not very familiar with pybind11, as I have just started looking into it. Any help on this is appreciated.

[saran-t]

The bindings weren't originally authored with C++ extensibility in mind, which makes it a bit difficult to pass things back into external C++ code.

The easiest way (but also hackiest way) to pass mjModel back to user-authored C++ code is to pass m._address to a function that accepts a uintptr_t as argument, reinterpret_cast it back to mjModel*, then pass it along to your actual function.

We can discuss options on how to do things more properly though, if you're interested!

[DavidPL1]

Indeed that sounds very hacky. And I'm very interested in discussing a proper solution.

I now have now split my loading process into two separate calls: initModelFromQueue which creates mjModel and mjData in C++ and loadModelAndData which replaces model and data of the environment class with the last initialized mjModel and mjData (this is very similar to what happens when Simulate is given a new model and data to display).

This now allows me to provide mjModel and mjData from python and directly issue a C++ call to loadModelAndData. The idea was to more or less copy what happens here https://github.com/deepmind/mujoco/blob/bbbf30ee47688bb03cefc7c3d6475b8283c0712f/python/mujoco/simulate.cc#L102-L117

But for the cast to mjModelWrapper and mjDataWrapper (and probably other wrappers later on) I would need to somehow include at least structs.h and I'm not sure if and how including headers from and linking a pybind shared library is possible.

[saran-t]

I mean, in some sense the hacky solution is the safest one too. Passing C++ objects across ABI boundaries is something that needs to be approached with caution. Passing raw C struct pointers around is safe. The wrapper classes are mostly just there to hold Python view objects, there's no real need for external C++ code to know about them.

Ideally we'd find a way to tell pybind11 to grab _address whenever binding to functions accepting mjModel*. But short of that doing manual casting achieved the same thing.

[DavidPL1]

Thanks for your input! I'll try to implement it like that.

Ideally we'd find a way to tell pybind11 to grab _address whenever binding to functions accepting mjModel*

I'll search for a way to automatically do that, but otherwise I think it would be just fine to manually override each function accepting respective pointers with a version where the cast happens before passing the casted pointers to the base class functions. In my case this should be limited to loading functions because otherwise I want the C++ program handling the objects. I only want to use python as an interactive way of triggering specific callbacks within C++ and maybe directly manipulating/reading contents of model and data.

(Communication over ROS requires specifying messages, services, and/or actions for each specific use case, whereas bindings provide full read and maybe write access more easily. That's the whole point why I'm doing this).

[DavidPL1]

I've tried casting manually, like you suggested, but there seems to be an issue with members of casted objects that are pointers.

I'm using this function that binds to mujoco_ros.test_conversion in python:

void test_conversion(py::object m, py::object d) {
    std::uintptr_t m_raw = m.attr("_address").cast<std::uintptr_t>();
    std::uintptr_t d_raw = d.attr("_address").cast<std::uintptr_t>();

    std::cout << std::showbase << std::hex;
    std::cout << "m_raw address: " << m_raw << std::endl;
    std::cout << "d_raw address: " << d_raw << std::endl;
    std::cout << std::dec;

    mjModel* m_cpp_ = reinterpret_cast<mjModel *>(m_raw);
    std::cout << "cpp model attributes: " << std::endl;
    std::cout << "\tnq: " << m_cpp_->nq << std::endl;
    std::cout << "\tnv: " << m_cpp_->nv << std::endl;
    std::cout << "\tnbody: " << m_cpp_->nbody << std::endl;
    std::cout << "\tngeom: " << m_cpp_->ngeom<< std::endl;
    std::cout << "\tqpos0: [";
    for (size_t i = 0; i < m_cpp_->nq; i++) {
        std::cout << m_cpp_->qpos0[i] << " ";
    }
    std::cout << "]" << std::endl;

    mjData* d_cpp_  = reinterpret_cast<mjData *>(d_raw);
    std::cout << "cpp data attributes: " << std::endl;
    std::cout << "\tnstack: " << d_cpp_->nstack << std::endl;
    std::cout << "\tnbuffer: " << d_cpp_->nbuffer << std::endl;
    std::cout << "\ttime: " << d_cpp_->time << std::endl;
    std::cout << "\tqpos: [";
    for (size_t i = 0; i < m_cpp_->nq; i++) {
        std::cout << d_cpp_->qpos[i] << " ";
    }
    std::cout << "]" << std::endl;
}

And tested the following python code:

import mujoco
import mujoco_ros

m = mujoco.MjModel.from_xml_path("mug.xml")
d = mujoco.MjData(m)

print(f"python m address: {hex(m._address)}")
print(f"python d address: {hex(d._address)}\n")

print(f"py model attributes:\n\tnq: {m.nq}\n\tnv: {m.nv}\n\tnbody: {m.nbody}\n\tngeom: {m.ngeom}\n\tqpos0: {m.qpos0}")

print(f"py data attributes: \n\tnstack: {d.nstack}\n\tnbuffer: {d.nbuffer}\n\tqpos[0]: {d.qpos[0]}\n\ttime: {d.time}\n")

mujoco_ros.test_conversion(m, d)
print("\n")

mujoco.mj_step(m, d)
print(f"py data attributes after step: \n\tnstack: {d.nstack}\n\tnbuffer: {d.nbuffer}\n\tqpos[0]: {d.qpos[0]}\n\ttime: {d.time}\n")
mujoco_ros.test_conversion(m, d)

And get the following output:

python m address: 0x2b42300
python d address: 0x3eb7280

py model attributes:
    nq: 7
    nv: 6
    nbody: 3
    ngeom: 36
    qpos0: [0. 0. 0. 1. 0. 0. 0.]
py data attributes:
    nstack: 1281240
    nbuffer: 3207376
    qpos[0]: 0.0
    time: 0.0

m_raw address: 0x2b42300
d_raw address: 0x3eb7280
cpp model attributes:
    nq: 7
    nv: 6
    nbody: 3
    ngeom: 0
    qpos0: [2.122e-314 0 0 0 0 0 0 ]
cpp data attributes:
    nstack: 1281240
    nbuffer: 3207376
    time: 3.25112e-316
    qpos: [0 0 0 0 0 0 0 ]

py data attributes after step:
    nstack: 1281240
    nbuffer: 3207376
    qpos[0]: -4.296789444571567e-22
    time: 0.002

m_raw address: 0x2b42300
d_raw address: 0x3eb7280
cpp model attributes:
    nq: 7
    nv: 6
    nbody: 3
    ngeom: 0
    qpos0: [2.122e-314 0 0 0 0 0 0 ]
cpp data attributes:
    nstack: 1281240
    nbuffer: 3207376
    time: 3.25112e-316
    qpos: [-1.0742e-16 -1.17204e-27 -9.81 -2.6475e-23 5.20977e-15 -3.29505e-24 0 ]

While the tested non-pointer attributes all match, the pointers to qpos, qpos0, and time seem to be garbage. Do you have any idea on what's the issue?

EDIT: Just noticed ngeom is also incorrect. And, the used "mug.xml" is this model.

[DavidPL1]

I just revisited this issue and realized that the version of the python bindings and the MuJoCo library version I was using mismatched. No wonder some pointers could not be interpreted correctly.

With matching versions everything seems to be correct.