refactor

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``

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:

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):

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

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(

mujoco_utils/arena.py → 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):

mujoco_utils/morphology.py → 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):

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

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
+        )