diff --git a/README.md b/README.md index 471b4f8..ff22c0b 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,51 @@ # mujoco-utils -Framework and utilities for interfacing with MuJoCo and MuJoCo-XLA. + +MuJoCo-utils provides a framework for implementing and interfacing with MuJoCo and MuJoCo-XLA simulation environments. +The main goal of this framework is to **unify** the development and interfaces of environments implemented in native +MuJoCo (MJC) and MuJoCo-XLA (MJX). + +## MJCFMorphology and MJCFArena + +* Both are parameterized/reconfigurable MJCF (MuJoCo-XML) generators that serve as input to an environment. +* MJCFMorphology defines the robot morphology. + * MJCFMorphologies are parameterized using + the [Framework for Parameterized Robot Specifications (FPRS)](https://github.com/Co-Evolve/fprs). + * MJCFMorphologies follow a modular design, dividing the robot into distinct parts. +* MJCFArena defines the arena in which the robot is places (i.e. all non-morphological structures). + * MJCFArena are reconfigurable via a `ArenaConfiguration` + * reconfigurable via a configuration + +## Unified MJC and MJX environment interface + +* Reconfigurable through an environment configuration +* Functional programming + * MJX requires the functional programming paradigm. To unify MJC and MJX, we thus apply this stronger coding + constraint to the MJC side as well. + * Simply said: functions are pure, and should not have side effects. + * Environments expect states and return states. +* Gymnasium-like interface + * Environments provide an observation and action space (automatically) + * Outside code interacts with the environment using the typical `step`, `reset`, `render` and `close` functions. + * Main difference with the original gymnasium interface: here the environment returns and expects states, which + encapsulates the original observations, rewards, info, etc. in a datastructure, together with the underlying + physics state. +* Differences between MJCEnv (i.e. a native MuJoCo environment) and MJXEnv (i.e. an MJX environment): + * development: MJXEnv's should be implemented in JAX, MJCEnv's do not necessarily require JAX. + * usage: MJXEnvState contains JAX arrays (and the additional mjx.Model and mjx.Data structures), the MJCEnvState + uses numpy arrays. +* Observations are implemented using Observables; these define retriever functions that extract the observation based on + the state datastructure and provide an easy way to implement observations. +* 'DualMuJocoEnvironment' provides the unification of MJC and MJX, and allows conditional environment creation based on + a backend string. +* Both MJC and MJX support human-mode and rgb_array rendering modes. + * Note: MJX rendering is slow due to the offloading of datastructures of the GPU + +## Examples + +For practical applications and demonstrations of MuJoCo-Utils, please refer to +the [Bio-inspired Robotics Benchmark](https://github.com/Co-Evolve/brb), +which employs this framework extensively. + +## Installation + +``python -m pip install git+https://github.com/Co-Evolve/mujoco-utils`` \ No newline at end of file diff --git a/mujoco_utils/environment/base.py b/mujoco_utils/environment/base.py index fd5aca4..102d264 100644 --- a/mujoco_utils/environment/base.py +++ b/mujoco_utils/environment/base.py @@ -12,9 +12,9 @@ from mujoco import mjx import mujoco_utils.environment.mjx_spaces as mjx_spaces -from mujoco_utils.arena import MJCFArena +from mujoco_utils.mjcf.arena import MJCFArena from mujoco_utils.environment.renderer import MujocoRenderer -from mujoco_utils.morphology import MJCFMorphology +from mujoco_utils.mjcf.morphology import MJCFMorphology class MuJoCoEnvironmentConfiguration: diff --git a/mujoco_utils/environment/dual.py b/mujoco_utils/environment/dual.py index 5c09752..27b49c2 100644 --- a/mujoco_utils/environment/dual.py +++ b/mujoco_utils/environment/dual.py @@ -6,11 +6,11 @@ import numpy as np from gymnasium.core import RenderFrame -from mujoco_utils.arena import MJCFArena +from mujoco_utils.mjcf.arena import MJCFArena from mujoco_utils.environment.base import BaseEnvState, BaseEnvironment, MuJoCoEnvironmentConfiguration, SpaceType from mujoco_utils.environment.mjc_env import MJCEnv from mujoco_utils.environment.mjx_env import MJXEnv -from mujoco_utils.morphology import MJCFMorphology +from mujoco_utils.mjcf.morphology import MJCFMorphology class DualMuJoCoEnvironment(BaseEnvironment): diff --git a/mujoco_utils/environment/mjc_env.py b/mujoco_utils/environment/mjc_env.py index 1f75793..fd4abee 100644 --- a/mujoco_utils/environment/mjc_env.py +++ b/mujoco_utils/environment/mjc_env.py @@ -14,10 +14,10 @@ from gymnasium.core import ActType, RenderFrame from gymnasium.vector.utils import batch_space -from mujoco_utils.arena import MJCFArena from mujoco_utils.environment.base import BaseEnvState, BaseEnvironment, BaseMuJoCoEnvironment, BaseObservable, \ MuJoCoEnvironmentConfiguration -from mujoco_utils.morphology import MJCFMorphology +from mujoco_utils.mjcf.arena import MJCFArena +from mujoco_utils.mjcf.morphology import MJCFMorphology @struct.dataclass diff --git a/mujoco_utils/environment/mjx_env.py b/mujoco_utils/environment/mjx_env.py index 7e85728..b714025 100644 --- a/mujoco_utils/environment/mjx_env.py +++ b/mujoco_utils/environment/mjx_env.py @@ -14,10 +14,10 @@ from mujoco import mjx import mujoco_utils.environment.mjx_spaces as mjx_spaces -from mujoco_utils.arena import MJCFArena from mujoco_utils.environment.base import BaseEnvState, BaseMuJoCoEnvironment, BaseObservable, \ MuJoCoEnvironmentConfiguration -from mujoco_utils.morphology import MJCFMorphology +from mujoco_utils.mjcf.arena import MJCFArena +from mujoco_utils.mjcf.morphology import MJCFMorphology def mjx_get_model( diff --git a/mujoco_utils/mjcf/__init__.py b/mujoco_utils/mjcf/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mujoco_utils/arena.py b/mujoco_utils/mjcf/arena.py similarity index 92% rename from mujoco_utils/arena.py rename to mujoco_utils/mjcf/arena.py index e115f39..2d5a9a8 100644 --- a/mujoco_utils/arena.py +++ b/mujoco_utils/mjcf/arena.py @@ -1,6 +1,6 @@ from abc import ABC -from mujoco_utils.mjcf_utils import MJCFRootComponent +from mujoco_utils.mjcf.component import MJCFRootComponent class ArenaConfiguration(ABC): diff --git a/mujoco_utils/mjcf_utils.py b/mujoco_utils/mjcf/component.py similarity index 100% rename from mujoco_utils/mjcf_utils.py rename to mujoco_utils/mjcf/component.py diff --git a/mujoco_utils/morphology.py b/mujoco_utils/mjcf/morphology.py similarity index 95% rename from mujoco_utils/morphology.py rename to mujoco_utils/mjcf/morphology.py index bf0c0f8..d948dc7 100644 --- a/mujoco_utils/morphology.py +++ b/mujoco_utils/mjcf/morphology.py @@ -8,7 +8,7 @@ from fprs.robot import Morphology from fprs.specification import MorphologySpecification -from mujoco_utils.mjcf_utils import MJCFRootComponent, MJCFSubComponent +from mujoco_utils.mjcf.component import MJCFRootComponent, MJCFSubComponent class MJCFMorphology(Morphology, MJCFRootComponent, ABC): diff --git a/requirements.txt b/requirements.txt index b18d5a1..e9c08e3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ chex==0.1.85 dm_control==1.0.16 flax==0.7.5 -fprs @ git+https://github.com/Co-Evolve/fprs@40ee809ca6ba15e66961d2153ee86e83550592c1 +fprs @ git+https://github.com/Co-Evolve/fprs@2512dcf827510740b7a7a83093fdbf3752107b16 glfw==2.6.4 gymnasium==0.29.1 imageio==2.33.1 diff --git a/setup.py b/setup.py index ef859ff..5281e24 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -from setuptools import setup, find_packages +from setuptools import find_packages, setup with open('README.md') as f: readme = f.read() @@ -10,12 +10,12 @@ required = f.read().splitlines() setup( - name='mujoco_utils', - version='0.2.0', - description='Utilities for interfacing with MuJoCo through dm_control.', - long_description=readme, - url='https://github.com/Co-Evolve/mujoco-utils', - license=license, - packages=find_packages(exclude=('tests', 'docs')), - install_requires=required -) + name='mujoco_utils', + version='1.0.0', + description='Framework and utilities for implementing and interfacing with MuJoCo and MuJoCo-XLA environments.', + long_description=readme, + url='https://github.com/Co-Evolve/mujoco-utils', + license=license, + packages=find_packages(exclude=('tests', 'docs')), + install_requires=required + )