Implementation of STAC using MJX. This is part of the VNL project.
stac-mjx relies on many prerequisites, therefore we suggest installing in a new conda environment, using the provided environment.yaml
:
[Local installation before package is officially published]
- Clone the repository
git clone https://github.com/talmolab/stac-mjx.git
andcd
into it - Create and activate the
stac-mjx-env
environment:
conda env create -f environment.yaml
conda activate stac-mjx-env
Our rendering functions support multiple backends: egl
, glfw
, and osmesa
. We show osmesa
setup as it supports headless rendering, which is common in remote/cluster setups. To set up (currently on supported on Linux), execute the following commands sequentially:
sudo apt-get install libglfw3 libglew2.0 libgl1-mesa-glx libosmesa6
conda install -c conda-forge glew
conda install -c conda-forge mesalib
conda install -c anaconda mesa-libgl-cos6-x86_64
conda install -c menpo glfw3
Finally, set the following environment variables, and reactivate the conda environment:
conda env config vars set MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa
conda deactivate && conda activate base
To ensure all of the above changes are encapsulated in your Jupyter kernel, create a new kernel with:
conda install ipykernel
python -m ipykernel install --user --name stac-mjx-env --display-name "Python (stac-mjx-env)"
-
Update the .yaml files in
config/
with the proper information (details WIP). -
Run stac-mjx with its basic api:
load_configs
for loading configs andrun_stac
for the keypoint registration. Below is an example script, found indemos/use_api.ipynb
.from stac_mjx import main from stac_mjx import utils from pathlib import Path # Set base path to the parent directory of your config files base_path = Path.cwd() stac_config_path = base_path / "demos/demo_stac.yaml" model_config_path = base_path / "configs/rodent.yaml" # Load configs cfg = main.load_configs(stac_config_path, model_config_path) # Load data data_path = base_path / cfg.paths.data_path kp_data = utils.load_data(data_path, utils.params) # Run stac fit_path, transform_path = main.run_stac(cfg, kp_data, base_path)
-
Render the resulting data using
mujoco_viz()
(example notebook found indemos/viz_usage.ipynb
):import os import mediapy as media from stac_mjx.viz import mujoco_viz from stac_mjx import main from stac_mjx import utils stac_config_path = "../configs/stac.yaml" model_config_path = "../configs/rodent.yaml" cfg = main.load_configs(stac_config_path, model_config_path) xml_path = "../models/rodent.xml" data_path = "../output.p" n_frames=250 save_path="../videos/direct_render.mp4" # Call mujoco_viz frames = mujoco_viz(data_path, xml_path, n_frames, save_path, start_frame=0) # Show the video in the notebook (it is also saved to the save_path) media.show_video(frames, fps=utils.params["RENDER_FPS"])
-
If the rendering is poor, it's likely that some hyperparameter tuning is necessary. (details WIP)