Conversation
- Add rollout() and sim() C++ functions with Python bindings - Add RolloutBackend class in judo.utils.mujoco_cpp module - Add rollout_backend config option to ControllerConfig to select between 'mujoco' (default) and 'mujoco_cpp' backends - Update README with instructions to build C++ extensions - Requires building C++ extensions with 'pixi run build' or manual CMake build
johnzhang3
requested review from
alberthli,
bhung-bdai,
jbruedigam-bdai,
pculbertson and
slecleach
as code owners
alberthli
left a comment
alberthli
left a comment
There was a problem hiding this comment.
High-level comments:
- need to resolve pip installation of the cpp package
- need to unify the rollout backend API
- need to add tests verifying that the parallel rollouts from your judo_cpp impl match the regular mujoco rollouts
- at least one person with more cpp competence than me should review this PR. not sure who, but please request it.
| @@ -1,3 +1,12 @@ | |||
| # v0.0.6 | |||
|
|
|||
| ## Added | |||
There was a problem hiding this comment.
tag yourself in this update to make sure you get credit!
| [tool.pixi.feature.cpp.dependencies] | ||
| cmake = "*" | ||
| ninja = "*" | ||
| pybind11 = "*" | ||
| cxx-compiler = "*" | ||
| c-compiler = "*" | ||
| eigen = "*" |
There was a problem hiding this comment.
nit: can we move this above [tool.pixi.environments]? it's weird that any dependency group definition comes below the environments.
There was a problem hiding this comment.
I don't believe Eigen is used here...
|
|
||
| # first time only | ||
| pre-commit install | ||
| pybind11-stubgen mujoco -o typings/ |
There was a problem hiding this comment.
after this block in the README, also explain how to clean the cpp build files using pixi run clean
|
|
||
| self.rollout_backend = RolloutBackend(num_threads=self.optimizer_cfg.num_rollouts, backend=rollout_backend) | ||
| # Determine backend: config override > parameter > "mujoco" | ||
| backend: Literal["mujoco", "mujoco_cpp"] = controller_config.rollout_backend or rollout_backend or "mujoco" |
There was a problem hiding this comment.
for readability reasons, I'd prefer if the logic for choosing the rollout backend wasn't so pythonic. a more explicit if/elif/else block would make this a lot easier to maintain
| if backend == "mujoco_cpp": | ||
| if CppRolloutBackend is None: | ||
| raise ImportError( | ||
| "judo_cpp module is not available. Please build the C++ extensions or use 'mujoco' backend." |
There was a problem hiding this comment.
you might want to explain what command to run explicitly here and point the user to the README
| Shutdown the persistent thread pool. | ||
|
|
||
| Call this function to clean up the persistent thread pool when done with rollouts. | ||
| The pool will be automatically recreated on the next call to persistent_cpp_rollout. |
There was a problem hiding this comment.
where is persistent_cpp_rollout defined?
| R"doc( | ||
| Run parallel MuJoCo rollouts. | ||
|
|
||
| Args: |
There was a problem hiding this comment.
nit: incorrect indentation for Args
| models: length-B list of mujoco._structs.MjModel | ||
| data: length-B list of mujoco._structs.MjData | ||
| x0: 2D array of shape (B, nq+nv), batched initial [qpos;qvel] | ||
| controls: 3D array of shape (B, horizon, nu), batched control inputs |
There was a problem hiding this comment.
nit: style-wise, we don't align text like this in python code/docstrings, so i'd prefer removing the extra tabbed whitespace in all these docstring descriptions (this goes for other stuff too)
| tuple of three np.ndarray: | ||
| states -> shape (B, horizon+1, nq+nv) - MuJoCo states (includes initial state) | ||
| sensors -> shape (B, horizon, nsensordata) - sensor data |
There was a problem hiding this comment.
instead of saying this, just write out the elements of the tuple following the same style as the usual docstring in the package. also, this says you return three, but two things are written here.
| R"doc( | ||
| Run a single MuJoCo simulation step. | ||
|
|
||
| Args: |
There was a problem hiding this comment.
nit: incorrect indentation
dyershov-bdai
left a comment
dyershov-bdai
left a comment
There was a problem hiding this comment.
It's a bit of a drive-by review. We need to talk a bit more about C++ best practices and arch. But this PR needs move work.
| // Thread pool for persistent thread management | ||
| class PersistentThreadPool { | ||
| public: | ||
| PersistentThreadPool(int num_threads); |
There was a problem hiding this comment.
Use rule-of-zero or rule-of-(five|seven).
| py::array_t<double> make_array(std::vector<double>& buf, int B, int T, int D); | ||
|
|
||
| // Thread pool for persistent thread management | ||
| class PersistentThreadPool { |
There was a problem hiding this comment.
You can use boost::thread_pool: https://www.boost.org/doc/libs/latest/doc/html/boost_asio/reference/thread_pool.html
| namespace py = pybind11; | ||
|
|
||
| // Utility function to create numpy arrays from C++ vectors | ||
| py::array_t<double> make_array(std::vector<double>& buf, int B, int T, int D); |
| const std::vector<const mjModel*>& models, | ||
| const std::vector<mjData*>& data, | ||
| const py::array_t<double>& x0, | ||
| const py::array_t<double>& controls |
There was a problem hiding this comment.
It's unsafe to hold references as member values. You may run into dangling reference if the referenced value moves out of scope.
| #include <unordered_map> | ||
|
|
||
| namespace py = pybind11; | ||
|
|
There was a problem hiding this comment.
Please define namespace. Do not put public API into global namespace.
| active_workers_--; | ||
| completed_tasks_++; |
There was a problem hiding this comment.
Prefer pre over post increment/decrement.
| }; | ||
|
|
||
| // Move buf onto the heap so we can own it in the capsule: | ||
| auto heap_buf = new std::vector<double>(std::move(buf)); |
|
|
||
| // ========== End Thread Pool Implementation ========== | ||
|
|
||
| py::array_t<double> make_array(std::vector<double>& buf, |
There was a problem hiding this comment.
If you wish to move, then change signature to
| py::array_t<double> make_array(std::vector<double>& buf, | |
| py::array_t<double> make_array(std::vector<double>&& buf, |
|
|
||
| auto controls_unchecked = controls.unchecked<3>(); | ||
|
|
||
| #pragma omp parallel for |
There was a problem hiding this comment.
If you are not going to use ThreadPoolManager, then why commit it in this PR?
| [tool.pixi.feature.cpp.dependencies] | ||
| cmake = "*" | ||
| ninja = "*" | ||
| pybind11 = "*" | ||
| cxx-compiler = "*" | ||
| c-compiler = "*" | ||
| eigen = "*" |
There was a problem hiding this comment.
I don't believe Eigen is used here...
4 participants
add cpp rollout features