MolmoSpaces

The molmo_spaces Python package is organized into several functional modules. Use these modules to build data generation pipelines, manage robot simulations, and evaluate policies:

• molmo_spaces: The top-level package containing core constants.
• configs: Classes for configuring experiments and individual components.
• controllers: Implementations of robot controllers.
• data_generation: Pipelines for generating data and the registry for their configurations.
• env: Base classes for defining simulation environments.
• evaluation: Infrastructure and pipelines for running evaluations.
• kinematics: Utilities and solvers for kinematic calculations.
• policy: Interfaces and wrappers for robot policies.
• renderer: Backends for rendering simulation views.
• resources: Bundled data files and scripts for generating user asset or grasp library indices.
• robots: Abstractions and models for different robots.
• tasks: Definitions for specific tasks and their associated reward functions.
• utils: Shared utility functions used across the package.

MolmoSpaces is a modular robotics simulation framework built on MuJoCo. Its architecture is organized into several functional domains:

• configs/: Centralized configuration management using a hierarchical system.
• controllers/: Low-level robot control interfaces (e.g., joint_pos.py, joint_vel.py).
• data_generation/: The pipeline for generating datasets, including the main.py entry point and pipeline.py for parallel rollouts.
• env/: Environment abstractions, including scene construction (arena/), sensor management, and the MuJoCo environment wrapper.
• evaluation/: Tools for benchmarking policies, including JSON runners and policy servers.
• grasp_generation/: Pipelines for processing meshes and generating grasps for rigid or articulated objects.
• kinematics/: Various kinematic solvers, including GPU-parallel and robot-specific implementations (e.g., franka_kinematics.py).
• planner/: Motion planning interfaces, including A* and CuRobo integration.
• policy/: Implementations of different control strategies, including learned policies, planning-based solvers, and dummy/random policies.
• renderer/: Rendering backends such as OpenGL and Filament.
• robots/: Abstract robot interfaces and specific implementations (e.g., franka.py, rby1.py).
• tasks/: Definitions of robot objectives (e.g., pick_task.py, nav_task.py) and task samplers.

The MolmoSpaces project documentation is distributed across several modules. Use the following index to find specific guides for asset generation, simulator usage, evaluation, and data processing:

Asset Generation

• Grasp generation: Instructions for generating grasps (../molmo_spaces/grasp_generation/README.md).
• House generation: Instructions for generating house environments (../molmo_spaces/housegen/README.md).

Assets Usage and Simulators

• MuJoCo: Asset and resource management (docs/assets.md).
• Isaac-sim: Usage within Isaac-sim (molmo_spaces_isaac/README.md).
• ManiSkill: Usage within ManiSkill (molmo_spaces_maniskill/README.md).

Evaluation

• Evaluation directions: General guidance on evaluation (../molmo_spaces/evaluation/README.md).
• Benchmarks & Comparisons: Detailed documentation for the MolmoSpaces benchmark and fair policy comparisons is available via external Google Doc links.

Data & Processing

• Data generation: Scripts and workflows for generating data (scripts/datagen/README.md).
• Data format: Specification of the data schema (docs/data_format.md).
• Data postprocessing: Procedures for processing generated data (docs/data_processing.md).

Development

• Code structure: Overview of the repository layout (docs/code_structure.md).
• Development: General development guidelines (docs/development.md).
• Tests: Information on running tests (mlspaces_tests/README.md).

The sensor system is an abstraction that converts raw simulator state into observation dictionaries consumed by policies and used for data generation.

Sensor Abstraction

A sensor is an object that returns a specific piece of an observation given the environment and task. Every sensor inherits from Sensor and includes:

• uuid: A unique string identifier used as the key in the observation dictionary and HDF5 files.
• observation_space: A gym.Space describing the output.
• is_dict: A boolean indicating if the output is a JSON-serializable dictionary.
• str_max_len: The padding length for the JSON byte buffer (used for dict sensors).
• get_observation(...): The method to retrieve the data.
• reset(): An optional method to clear internal sensor state (e.g., for sensors that cache previous poses).

Sensor Output Types

1. Plain array sensors (is_dict = False): Return a np.ndarray (e.g., Camera RGB, depth, TCP pose).
2. Dict sensors (is_dict = True): Return a Python dict. These are serialized via json.dumps, UTF-8 encoded, and packed into a fixed-length np.uint8 buffer of length str_max_len. The observation_space is a Box(0, 255, (str_max_len,), uint8).

class Sensor(ABC):
    uuid: str                       # unique identifier, used as the obs dict key
    observation_space: gym.Space    # gymnasium space describing the output
    is_dict: bool = False           # if True, output is a dict that will be JSON-encoded
    str_max_len: int = 2000         # padding length for the JSON byte buffer

    @abstractmethod
    def get_observation(self, env, task, batch_index: int = 0, ...): ...

    def reset(self) -> None: ...    # optional, override if the sensor has state

A MoveGroup abstracts a collection of joints and actuators. It allows you to interact with a robot part (like an arm or gripper) without needing to know the underlying MuJoCo joint names or actuator mappings.

Key Interface

• State: Access/set joint_pos, joint_vel, and ctrl as numpy arrays.
• Limits: Check joint_pos_limits and ctrl_limits.
• Frames: Access transforms via leaf_frame_to_world, root_frame_to_world, and leaf_frame_to_root.
• Control: Use get_jacobian() to get the Jacobian mapping joint velocities to the spatial velocity of the leaf frame, or noop_ctrl for no-op control.

Common Specializations

• SimplyActuatedMoveGroup: Used when there is a 1:1 mapping between joints and actuators (n_joints == pos_dim == vel_dim == n_actuators).
• GripperGroup: Adds gripper-specific methods like set_gripper_ctrl_open, is_open, and inter_finger_dist.
• RobotBaseGroup: Represents the robot's pose in the world (e.g., MocapRobotBaseGroup for a fixed tabletop base or FreeJointRobotBaseGroup for a 6-DoF base).

The data generation process follows a structured flow to produce datasets through parallel rollouts:

1. Entry Point: Execution starts via molmo_spaces.data_generation.main, which loads an experiment configuration.
2. Configuration: The system uses MlSpacesExpConfig to define the experiment, which includes configurations for the Task Sampler, Robot, Policy, Camera, and fixed task parameters.
3. Pipeline Execution: The ParallelRolloutRunner manages multi-process execution by:
   • Creating task samplers for each worker.
   • Sampling tasks based on defined parameters.
   • Initializing the necessary policies and environments.
   • Running rollouts and collecting the resulting data.
4. Episode Lifecycle: For every episode, the sequence is: Task Sampling $\rightarrow$ Policy Initialization $\rightarrow$ Episode Rollout $\rightarrow$ Success Evaluation.

The molmo_spaces JSON evaluation pipeline follows a structured lifecycle from CLI/programmatic entry to result collection. The process is managed primarily through run_evaluation() in molmo_spaces/evaluation/eval_main.py and executed by the JsonEvalRunner class.

High-level Execution Flow

1. Initialization: run_evaluation() resolves the evaluation config class, loads benchmark episodes from JSON, and resolves the task_horizon (prioritizing CLI overrides over the benchmark's task_horizon_sec).
2. Configuration: An evaluation config is instantiated via create_eval_config(), which enforces evaluation-mode flags (e.g., seed=42, no action noise, no datagen profiler). CLI overrides for camera_config, camera_names, and light intensity are then applied.
3. Runner Setup: JsonEvalRunner is initialized. It loads episodes from the benchmark_dir, handles truncation (via max_episodes or episode_idx), and derives task_sampler_config.house_inds and samples_per_house. It inherits from ParallelRolloutRunner to set up worker processes.
4. Execution: JsonEvalRunner.run() dispatches work items to workers. Each worker executes process_single_house(), which iterates through EpisodeSpec objects, samples tasks using JsonEvalTaskSampler, and runs rollouts at the specified policy_dt_ms.
5. Output: Trajectories (trajectories.h5) and per-episode artifacts are written to the output_dir. Finally, collect_episode_results() generates an EvaluationResults object containing success_count, total_count, output_dir, episode_results, and exp_config.

CLI args / programmatic call
        │
        ▼
run_evaluation()  (molmo_spaces/evaluation/eval_main.py)
  ├─ resolve eval config class (registry name or "module:Class")
  ├─ load benchmark episodes from JSON
  ├─ resolve task_horizon (CLI override > benchmark task_horizon_sec)
  ├─ create_eval_config(): instantiate eval config, force eval-mode flags
  │     (no action noise, no datagen profiler, seed=42, output_dir, ...)
  ├─ apply CLI overrides (camera_config, camera_names, light intensity)
  ├─ JsonEvalRunner.patch_config(): attach EvalRuntimeParams
  └─ JsonEvalRunner.adjust_robot(): wire robot_eval_override (if any)
        │
        ▼
JsonEvalRunner.__init__()
  ├─ load_all_episodes(benchmark_dir)
  ├─ truncate to max_episodes / single episode_idx
  ├─ derive task_sampler_config.house_inds + samples_per_house
  └─ super().__init__()  -> ParallelRolloutRunner sets up workers
        │
        ▼
JsonEvalRunner.run() -> ParallelRolloutRunner.run()
  └─ for each (house_id, batch) work item, dispatch to workers
        │
        ▼
ParallelRolloutRunner.process_single_house()  (per worker)
  ├─ load_episodes_for_house()       # JsonEvalRunner override
  ├─ for each EpisodeSpec:
  │     prepare_episode_config()
  │     get_episode_task_sampler()    -> JsonEvalTaskSampler(exp_config, ep)
  │     sample_task_from_spec()       -> task = sampler.sample_task(...)
  │     run_single_rollout()          -> step env at policy_dt_ms
  │     should_close_episode_task_sampler()  (True for JSON eval)
  └─ write trajectories.h5 + per-episode artifacts under output_dir
        │
        ▼
collect_episode_results() + (optional) wandb logging
        │
        ▼
EvaluationResults (success_count, total_count, output_dir, episode_results, exp_config)