data_generation

data_generation

Modules:

Name	Description
config	Data generation configuration files.
config_registry	Configuration registry for data generation experiments.
main
pipeline

config

Data generation configuration files.

All config classes in this directory will be automatically discovered and registered when the main data generation script runs.

Modules:

Name	Description
benchmarks_datagen_configs	Sepeare file containing imported benchmark tasks.
door_opening_configs
nav_to_obj_configs	Data generation configs for RBY1 navigation to object tasks.
object_manipulation_datagen_configs	Data generation configs for Franka move-to-pose tasks.

benchmarks_datagen_configs

Sepeare file containing imported benchmark tasks.

Why: Some old benchmarks ~18 DEC 2025 were generated with these, and because we pickle the tasks they need to be accessible in the same location. When no longer using the old benchmarks this file can be removed.

Moving benchmarks configs to this file caused some complications with imports, so we undid that change.

door_opening_configs

Classes:

Name	Description
DoorOpeningDataGenConfig	All-ProcTHOR variant for RBY1 door opening dataset generation.
DoorOpeningDebugConfig	Debug config for door opening dataset generation.

DoorOpeningDataGenConfig

Bases: MlSpacesExpConfig

All-ProcTHOR variant for RBY1 door opening dataset generation.

Iterates through multiple houses from the ProcTHOR dataset for large-scale data generation. This is the main config - use DoorOpeningDebugConfig for single-scene testing.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize policy config after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	RBY1GoProD455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	DoorOpeningPolicyConfig | None
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	RBY1MConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	DoorOpeningTaskConfig
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int
task_sampler_config	DoorOpeningTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: RBY1GoProD455CameraSystem = RBY1GoProD455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 20.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / 'experiment_output'

policy_config class-attribute instance-attribute

policy_config: DoorOpeningPolicyConfig | None = None

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 100.0

profile class-attribute instance-attribute

profile: bool = False

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: RBY1MConfig = RBY1MConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 4.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: DoorOpeningTaskConfig = DoorOpeningTaskConfig(task_cls=DoorOpeningTask)

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int = 400

task_sampler_config class-attribute instance-attribute

task_sampler_config: DoorOpeningTaskSamplerConfig = DoorOpeningTaskSamplerConfig(task_sampler_class=DoorOpeningTaskSampler)

task_type class-attribute instance-attribute

task_type: str = 'door_open'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'camera': 'robot_0/camera_follower'}

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	BaseRobotConfig | None
task_cls_str	str | None
task_config	AllTaskConfigs | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: AllTaskConfigs | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize policy config after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/door_opening_configs.py

def model_post_init(self, __context) -> None:
    """Initialize policy config after Pydantic model initialization"""
    super().model_post_init(__context)
    # Set up policy config with dynamically computed planner configs
    # Skip if no GPU available (e.g., when launching jobs from manager)
    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

DoorOpeningDebugConfig

Bases: DoorOpeningDataGenConfig

Debug config for door opening dataset generation.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize policy config after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
tag	A string describing the experiment.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	RBY1GoProD455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	DoorOpeningPolicyConfig | None
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	RBY1MConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
task_config	DoorOpeningTaskConfig
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int
task_sampler_config	DoorOpeningTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: RBY1GoProD455CameraSystem = RBY1GoProD455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 20.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = False

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / 'experiment_output' / 'door_opening_debug'

policy_config class-attribute instance-attribute

policy_config: DoorOpeningPolicyConfig | None = None

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 100.0

profile class-attribute instance-attribute

profile: bool = False

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: RBY1MConfig = RBY1MConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = 83067780

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 4.0

task_config class-attribute instance-attribute

task_config: DoorOpeningTaskConfig = DoorOpeningTaskConfig(task_cls=DoorOpeningTask)

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int = 1000

task_sampler_config class-attribute instance-attribute

task_sampler_config: DoorOpeningTaskSamplerConfig = DoorOpeningTaskSamplerConfig(task_sampler_class=DoorOpeningTaskSampler, samples_per_house=1, house_inds=[22])

task_type class-attribute instance-attribute

task_type: str = 'door_open'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = True

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'camera': 'robot_0/camera_follower'}

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	BaseRobotConfig | None
task_cls_str	str | None
task_config	AllTaskConfigs | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: AllTaskConfigs | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize policy config after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/door_opening_configs.py

def model_post_init(self, __context) -> None:
    """Initialize policy config after Pydantic model initialization"""
    super().model_post_init(__context)
    # Set up policy config with dynamically computed planner configs
    # Skip if no GPU available (e.g., when launching jobs from manager)
    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

tag

tag() -> str

A string describing the experiment.

Source code in molmo_spaces/data_generation/config/door_opening_configs.py

def tag(self) -> str:
    return "rby1_door_opening_debug"

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

nav_to_obj_configs

Data generation configs for RBY1 navigation to object tasks.

These configs subclass from NavToObjBaseConfig and are registered for use in the data generation pipeline.

Classes:

Name	Description
NavToObjDataGenConfig	Data generation config for RBY1 navigation to object tasks.

NavToObjDataGenConfig

Bases: NavToObjBaseConfig

Data generation config for RBY1 navigation to object tasks.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	RBY1MjcfCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_threads	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
record_videos	bool
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	NavToObjTaskConfig
task_config_preset	NavToObjTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int
task_sampler_config	NavToObjTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: RBY1MjcfCameraSystem = RBY1MjcfCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_threads class-attribute instance-attribute

num_threads: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'nav_to_obj_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = AStarNavToObjPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 200.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

record_videos class-attribute instance-attribute

record_videos: bool = False

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = RBY1Config()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: NavToObjTaskConfig = NavToObjTaskConfig(task_cls=NavToObjTask)

task_config_preset class-attribute instance-attribute

task_config_preset: NavToObjTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: NavToObjTaskSamplerConfig = NavToObjTaskSamplerConfig(task_sampler_class=NavToObjTaskSampler, pickup_types=None, robot_safety_radius=0.35, robot_object_z_offset=0.1, base_pose_sampling_radius_range=(4.0, 20.0), face_target=False, max_robot_placement_attempts=10, filter_for_successful_trajectories=True)

task_type class-attribute instance-attribute

task_type: str = 'nav_to_obj'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': array([0.0, 0.0, 0.5])}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	RBY1MjcfCameraSystem | None
robot_config	RBY1Config | None
task_cls_str	str | None
task_config	NavToObjTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: RBY1MjcfCameraSystem | None = None

robot_config class-attribute instance-attribute

robot_config: RBY1Config | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: NavToObjTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/nav_to_obj_configs.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    if not self.task_sampler_config.house_inds:
        self.task_sampler_config.house_inds = list(range(4))

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

object_manipulation_datagen_configs

Data generation configs for Franka move-to-pose tasks.

These configs subclass from the base_pick_config and are registered for use in the data generation pipeline.

Classes:

Name	Description
FrankaCloseDataGenConfig	Data generation config for Franka open task.
FrankaCloseHardBench	Data generation config for Franka open task.
FrankaOpenDataGenConfig	Data generation config for Franka open task.
FrankaOpenHardBench	Data generation config for Franka open task.
FrankaPickAndPlaceColorDataGenConfig
FrankaPickAndPlaceColorDroidDataGenConfig
FrankaPickAndPlaceColorOmniCamConfig	Data generation config for Franka pick task with Omni-directional cameras.
FrankaPickAndPlaceDataGenConfig
FrankaPickAndPlaceDroidDataGenConfig
FrankaPickAndPlaceEasyDataGenConfig
FrankaPickAndPlaceGoProD405D455DataGenConfig
FrankaPickAndPlaceGoProD405D455DataGenConfigDebug	Data generation config for Franka pick and place task with GoPro D405 cameras - deterministic version.
FrankaPickAndPlaceNextToDataGenConfig
FrankaPickAndPlaceNextToDroidDataGenConfig
FrankaPickAndPlaceNextToOmniCamConfig	Data generation config for Franka pick task with Omni-directional cameras.
FrankaPickAndPlaceOmniCamConfig	Data generation config for Franka pick task with Omni-directional cameras.
FrankaPickBatchTestConfig	Minimal config to test local batch-based work distribution.
FrankaPickDroidBench
FrankaPickDroidDataGenConfig	Data generation config for Franka pick task with DROID-style fixed cameras.
FrankaPickDroidMiniBench
FrankaPickGoProD405D455DataGenConfig	Data generation config for Franka pick task with GoPro D405 cameras.
FrankaPickHardBench
FrankaPickOmniCamAblationConfig	Data generation config for Franka pick task with Omni-directional cameras.
FrankaPickOmniCamConfig	Data generation config for Franka pick task with Omni-directional cameras.
FrankaPickRandomizedDataGenConfig	Data generation config for Franka pick task with randomized exocentric cameras.
FrankaPickandPlaceColorDroidBench	Data generation config for Franka pick task with DROID-style fixed cameras.
FrankaPickandPlaceColorHardBench
FrankaPickandPlaceDroidBench
FrankaPickandPlaceDroidMiniBench
FrankaPickandPlaceHardBench
FrankaPickandPlaceNextToDroidBench
FrankaPickandPlaceNextToHardBench
RBY1OpenDataGenConfig
RBY1PickAndPlaceDataGenConfig
RBY1PickDataGenConfig
RUMPickAndPlaceMultiDataGenConfig
RUMPickDataGenConfig

FrankaCloseDataGenConfig

Bases: ClosingBaseConfig

Data generation config for Franka open task.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	OpeningTaskConfig
task_config_preset	OpeningTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	OpenTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'close_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = OpenClosePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'ithor'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: OpeningTaskConfig = OpeningTaskConfig(task_cls=OpeningTask, task_success_threshold=0.85, joint_index=0, any_inst_of_category=False)

task_config_preset class-attribute instance-attribute

task_config_preset: OpeningTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 200

task_sampler_config class-attribute instance-attribute

task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(task_sampler_class=OpenTaskSampler, target_initial_state_open_percentage=0.5)

task_type class-attribute instance-attribute

task_type: str = 'close'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaCloseHardBench

Bases: ClosingBaseConfig

Data generation config for Franka open task.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	OpeningTaskConfig
task_config_preset	OpeningTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	OpenTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'close_bench'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = OpenClosePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig(init_qpos_noise_range={'arm': [0.26] * 6 + [pi / 2]})

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'ithor'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: OpeningTaskConfig = OpeningTaskConfig(task_cls=OpeningTask, task_success_threshold=0.85, joint_index=0, any_inst_of_category=False)

task_config_preset class-attribute instance-attribute

task_config_preset: OpeningTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 200

task_sampler_config class-attribute instance-attribute

task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(task_sampler_class=OpenTaskSampler, target_initial_state_open_percentage=0.5, robot_object_z_offset_random_min=-0.25, robot_object_z_offset_random_max=0.25, robot_placement_rotation_range_rad=0.52)

task_type class-attribute instance-attribute

task_type: str = 'close'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaOpenDataGenConfig

Bases: OpeningBaseConfig

Data generation config for Franka open task.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	OpeningTaskConfig
task_config_preset	OpeningTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	OpenTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'open_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = OpenClosePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'ithor'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: OpeningTaskConfig = OpeningTaskConfig(task_cls=OpeningTask, task_success_threshold=0.15, joint_index=0, any_inst_of_category=True)

task_config_preset class-attribute instance-attribute

task_config_preset: OpeningTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 200

task_sampler_config class-attribute instance-attribute

task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(task_sampler_class=OpenTaskSampler, target_initial_state_open_percentage=0)

task_type class-attribute instance-attribute

task_type: str = 'open'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaOpenHardBench

Bases: OpeningBaseConfig

Data generation config for Franka open task.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	OpeningTaskConfig
task_config_preset	OpeningTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	OpenTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'open_bench'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = OpenClosePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig(init_qpos_noise_range={'arm': [0.26] * 6 + [pi / 2]})

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'ithor'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: OpeningTaskConfig = OpeningTaskConfig(task_cls=OpeningTask, task_success_threshold=0.15, joint_index=0, any_inst_of_category=True)

task_config_preset class-attribute instance-attribute

task_config_preset: OpeningTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 200

task_sampler_config class-attribute instance-attribute

task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(task_sampler_class=OpenTaskSampler, target_initial_state_open_percentage=0, robot_object_z_offset_random_min=-0.25, robot_object_z_offset_random_max=0.25, robot_placement_rotation_range_rad=0.52)

task_type class-attribute instance-attribute

task_type: str = 'open'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceColorDataGenConfig

Bases: PickAndPlaceColorDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaRandomizedD405D455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceColorPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	FrankaRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceColorTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceColorTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaRandomizedD405D455CameraSystem = FrankaRandomizedD405D455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_colors_randomized_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceColorPlannerPolicyConfig = PickAndPlaceColorPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceColorTaskConfig = PickAndPlaceColorTaskConfig(task_cls=PickAndPlaceColorTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceColorTaskSamplerConfig = PickAndPlaceColorTaskSamplerConfig(task_sampler_class=PickAndPlaceColorTaskSampler, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_color'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceColorDroidDataGenConfig

Bases: PickAndPlaceColorDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceColorPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	FrankaRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceColorTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceColorTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_colors_droid_randomized_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceColorPlannerPolicyConfig = PickAndPlaceColorPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceColorTaskConfig = PickAndPlaceColorTaskConfig(task_cls=PickAndPlaceColorTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceColorTaskSamplerConfig = PickAndPlaceColorTaskSamplerConfig(task_sampler_class=PickAndPlaceColorTaskSampler, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_color'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceColorOmniCamConfig

Bases: PickAndPlaceColorDataGenConfig

Data generation config for Franka pick task with Omni-directional cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceColorPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceColorTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceColorTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_color_omni_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceColorPlannerPolicyConfig = PickAndPlaceColorPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceColorTaskConfig = PickAndPlaceColorTaskConfig(task_cls=PickAndPlaceColorTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceColorTaskSamplerConfig = PickAndPlaceColorTaskSamplerConfig(task_sampler_class=PickAndPlaceColorTaskSampler, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_color'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceDataGenConfig

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaRandomizedDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaRandomizedDroidCameraSystem = FrankaRandomizedDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_randomized_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=[], samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceDroidDataGenConfig

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_droid_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=[], samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceEasyDataGenConfig

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaEasyRandomizedDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaEasyRandomizedDroidCameraSystem = FrankaEasyRandomizedDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_randomized_easy_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=[], samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceGoProD405D455DataGenConfig

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaGoProD405D455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaGoProD405D455CameraSystem = FrankaGoProD405D455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_go_pro_d405_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=[], samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceGoProD405D455DataGenConfigDebug

Bases: FrankaPickAndPlaceDroidDataGenConfig

Data generation config for Franka pick and place task with GoPro D405 cameras - deterministic version.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaGoProD405D455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaGoProD405D455CameraSystem = FrankaGoProD405D455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = False

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'debug'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_go_pro_d405_v1_debug'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int = 100

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, samples_per_house=10, max_tasks=100, house_inds=[2])

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceNextToDataGenConfig

Bases: PickAndPlaceNextToDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaRandomizedD405D455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceNextToPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceNextToTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceNextToTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaRandomizedD405D455CameraSystem = FrankaRandomizedD405D455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_next_to_randomized_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceNextToPlannerPolicyConfig = PickAndPlaceNextToPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceNextToTaskConfig = PickAndPlaceNextToTaskConfig(task_cls=PickAndPlaceNextToTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceNextToTaskSamplerConfig = PickAndPlaceNextToTaskSamplerConfig(task_sampler_class=PickAndPlaceNextToTaskSampler, pickup_types=PICK_AND_PLACE_OBJECTS, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_next_to'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceNextToDroidDataGenConfig

Bases: PickAndPlaceNextToDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceNextToPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceNextToTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceNextToTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_next_to_droid_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceNextToPlannerPolicyConfig = PickAndPlaceNextToPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceNextToTaskConfig = PickAndPlaceNextToTaskConfig(task_cls=PickAndPlaceNextToTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceNextToTaskSamplerConfig = PickAndPlaceNextToTaskSamplerConfig(task_sampler_class=PickAndPlaceNextToTaskSampler, pickup_types=PICK_AND_PLACE_OBJECTS, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_next_to'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceNextToOmniCamConfig

Bases: PickAndPlaceNextToDataGenConfig

Data generation config for Franka pick task with Omni-directional cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceNextToPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceNextToTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceNextToTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_next_to_omni_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceNextToPlannerPolicyConfig = PickAndPlaceNextToPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceNextToTaskConfig = PickAndPlaceNextToTaskConfig(task_cls=PickAndPlaceNextToTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceNextToTaskSamplerConfig = PickAndPlaceNextToTaskSamplerConfig(task_sampler_class=PickAndPlaceNextToTaskSampler, pickup_types=PICK_AND_PLACE_OBJECTS, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_next_to'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickAndPlaceOmniCamConfig

Bases: PickAndPlaceDataGenConfig

Data generation config for Franka pick task with Omni-directional cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_and_place_omni_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=[], samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickBatchTestConfig

Bases: PickBaseConfig

Minimal config to test local batch-based work distribution.

Single house, 12 episodes split into 3 batches of 4, processed by 2 workers. Output: house_7/trajectories_batch_{1,2,3}_of_3.h5

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'debug'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 2

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_batch_test'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler, house_inds=[7], samples_per_house=12, episodes_per_batch=4)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickDroidBench

Bases: PickBaseConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_obja_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler, samples_per_house=40, house_inds=list(range(101)))

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickDroidDataGenConfig

Bases: PickBaseConfig

Data generation config for Franka pick task with DROID-style fixed cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_droid_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickDroidMiniBench

Bases: PickBaseConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_droid_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler, samples_per_house=40, house_inds=list(range(101)))

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickGoProD405D455DataGenConfig

Bases: PickBaseConfig

Data generation config for Franka pick task with GoPro D405 cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaGoProD405D455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaGoProD405D455CameraSystem = FrankaGoProD405D455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 4

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_go_pro_d405_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int = 150

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickHardBench

Bases: PickBaseConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_hard_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig(init_qpos_noise_range={'arm': [0.26] * 6 + [pi / 2]})

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler, robot_object_z_offset_random_min=-0.25, robot_object_z_offset_random_max=0.25, robot_placement_rotation_range_rad=0.52)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickOmniCamAblationConfig

Bases: FrankaPickOmniCamConfig

Data generation config for Franka pick task with Omni-directional cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_omni_v1_cam_ablation'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler, added_pickup_objects=None)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py

def model_post_init(self, __context) -> None:
    super().model_post_init(__context)
    if self.task_sampler_config.added_pickup_objects is None:
        self.task_sampler_config.added_pickup_objects = self._get_valid_pickupable_obja_uids()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickOmniCamConfig

Bases: PickBaseConfig

Data generation config for Franka pick task with Omni-directional cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_omni_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickRandomizedDataGenConfig

Bases: PickBaseConfig

Data generation config for Franka pick task with randomized exocentric cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaRandomizedD405D455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaRandomizedD405D455CameraSystem = FrankaRandomizedD405D455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pick_randomized_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceColorDroidBench

Bases: PickAndPlaceColorDataGenConfig

Data generation config for Franka pick task with DROID-style fixed cameras.

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceColorPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceColorTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceColorTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_color_obja_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceColorPlannerPolicyConfig = PickAndPlaceColorPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceColorTaskConfig = PickAndPlaceColorTaskConfig(task_cls=PickAndPlaceColorTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceColorTaskSamplerConfig = PickAndPlaceColorTaskSamplerConfig(task_sampler_class=PickAndPlaceColorTaskSampler, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_color'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceColorHardBench

Bases: PickAndPlaceColorDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceColorPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceColorTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceColorTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_color_hard_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceColorPlannerPolicyConfig = PickAndPlaceColorPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig(init_qpos_noise_range={'arm': [0.26] * 6 + [pi / 2]})

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceColorTaskConfig = PickAndPlaceColorTaskConfig(task_cls=PickAndPlaceColorTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceColorTaskSamplerConfig = PickAndPlaceColorTaskSamplerConfig(task_sampler_class=PickAndPlaceColorTaskSampler, robot_object_z_offset_random_min=-0.25, robot_object_z_offset_random_max=0.25, robot_placement_rotation_range_rad=0.52)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_color'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceDroidBench

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_obja_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=PICK_AND_PLACE_OBJECTS, samples_per_house=40, house_inds=list(range(101)))

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceDroidMiniBench

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_droid_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=PICK_AND_PLACE_OBJECTS, samples_per_house=40, house_inds=list(range(101)))

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceHardBench

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_hard_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig(init_qpos_noise_range={'arm': [0.26] * 6 + [pi / 2]})

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, robot_object_z_offset_random_min=-0.25, robot_object_z_offset_random_max=0.25, robot_placement_rotation_range_rad=0.52)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceNextToDroidBench

Bases: PickAndPlaceNextToDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceNextToPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceNextToTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceNextToTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaDroidCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_next_to_obja_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceNextToPlannerPolicyConfig = PickAndPlaceNextToPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceNextToTaskConfig = PickAndPlaceNextToTaskConfig(task_cls=PickAndPlaceNextToTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceNextToTaskSamplerConfig = PickAndPlaceNextToTaskSamplerConfig(task_sampler_class=PickAndPlaceNextToTaskSampler, pickup_types=PICK_AND_PLACE_OBJECTS, samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_next_to'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

FrankaPickandPlaceNextToHardBench

Bases: PickAndPlaceNextToDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaOmniPurposeCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlaceNextToPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	BaseRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceNextToTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaOmniPurposeCameraSystem = FrankaOmniPurposeCameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'val'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'benchmark' / 'pick_and_place_next_to_hard_v1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlaceNextToPlannerPolicyConfig = PickAndPlaceNextToPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: BaseRobotConfig = FrankaRobotConfig(init_qpos_noise_range={'arm': [0.26] * 6 + [pi / 2]})

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceNextToTaskConfig = PickAndPlaceNextToTaskConfig(task_cls=PickAndPlaceNextToTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceNextToTaskSamplerConfig(task_sampler_class=PickAndPlaceNextToTaskSampler, robot_object_z_offset_random_min=-0.25, robot_object_z_offset_random_max=0.25, robot_placement_rotation_range_rad=0.52)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place_next_to'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

RBY1OpenDataGenConfig

Bases: OpeningBaseConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	RBY1GoProD455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	BasePolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	RBY1MOpenCloseConfig
scene_dataset	str
seed	int
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	OpeningTaskConfig
task_config_preset	OpeningTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	OpenTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: RBY1GoProD455CameraSystem = RBY1GoProD455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 20.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'rby1_open_v1'

policy_config class-attribute instance-attribute

policy_config: BasePolicyConfig = CuroboOpenClosePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 100.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: RBY1MOpenCloseConfig = RBY1MOpenCloseConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'ithor'

seed class-attribute instance-attribute

seed: int = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 4.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: OpeningTaskConfig = OpeningTaskConfig(task_cls=OpeningTask, task_success_threshold=0.15, joint_index=0, any_inst_of_category=True)

task_config_preset class-attribute instance-attribute

task_config_preset: OpeningTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 200

task_sampler_config class-attribute instance-attribute

task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(task_sampler_class=OpenTaskSampler, target_initial_state_open_percentage=0, robot_safety_radius=0.2, base_pose_sampling_radius_range=(0.3, 1.0))

task_type class-attribute instance-attribute

task_type: str = 'open'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'mujoco-thor-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py

def model_post_init(self, __context) -> None:
    super().model_post_init(__context)
    self.policy_config = self._init_policy_config()
    self.task_config.task_success_threshold = 0.67
    self.task_sampler_config.randomize_textures = True

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

RBY1PickAndPlaceDataGenConfig

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	RBY1GoProD455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	CuroboPickAndPlacePlannerPolicyConfig | None
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	RBY1MConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: RBY1GoProD455CameraSystem = RBY1GoProD455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 20.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'rby1_pick_and_place_v1'

policy_config class-attribute instance-attribute

policy_config: CuroboPickAndPlacePlannerPolicyConfig | None = None

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 100.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: RBY1MConfig = RBY1MConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 4.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 400

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceTaskSampler, pickup_types=[], samples_per_house=20)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'camera': 'robot_0/camera_follower'}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'mujoco-thor-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py

def model_post_init(self, __context) -> None:
    super().model_post_init(__context)
    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise
    self.robot_config.init_qpos["head"][1] = 0.6
    self.task_sampler_config.robot_safety_radius = 0.35
    self.task_sampler_config.max_robot_to_obj_dist = 0.5
    self.task_sampler_config.object_placement_radius_range = (0.1, 0.5)
    self.task_sampler_config.min_object_to_receptacle_dist = 0.05
    self.task_sampler_config.max_robot_to_place_receptacle_dist = 0.5

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

RBY1PickDataGenConfig

Bases: PickBaseConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	RBY1GoProD455CameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	CuroboPickAndPlacePlannerPolicyConfig | None
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	RBY1MConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: RBY1GoProD455CameraSystem = RBY1GoProD455CameraSystem()

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 20.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'rby1_pick_v1'

policy_config class-attribute instance-attribute

policy_config: CuroboPickAndPlacePlannerPolicyConfig | None = None

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 100.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: RBY1MConfig = RBY1MConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 4.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 400

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig(task_sampler_class=PickTaskSampler)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'camera': 'robot_0/camera_follower'}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py

def model_post_init(self, __context) -> None:
    super().model_post_init(__context)
    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise
    self.robot_config.init_qpos["head"][1] = 0.6
    self.task_sampler_config.robot_safety_radius = 0.35
    self.task_sampler_config.max_robot_to_obj_dist = 0.5
    self.task_sampler_config.object_placement_radius_range = (0.1, 0.5)

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

RUMPickAndPlaceMultiDataGenConfig

Bases: PickAndPlaceDataGenConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickAndPlacePlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	FloatingRUMRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickAndPlaceTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	PickAndPlaceTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaRandomizedD405D455CameraSystem(img_resolution=(960, 720), visibility_constraints=None, allow_relaxed_constraints=True)

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'pnpmulti_V1'

policy_config class-attribute instance-attribute

policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: FloatingRUMRobotConfig = FloatingRUMRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'procthor-10k'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(task_sampler_class=PickAndPlaceMultiTaskSampler, pickup_types=None, samples_per_house=20, house_inds=[7], robot_object_z_offset=0.2, check_robot_placement_visibility=False)

task_type class-attribute instance-attribute

task_type: str = 'pick_and_place'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str = 'mujoco-thor-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickAndPlaceTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickAndPlaceTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

RUMPickDataGenConfig

Bases: PickBaseConfig

Classes:

Name	Description
Config
SavedEpisode	Config informationd describing a sinlge episode

Methods:

Name	Description
freeze_task_config	Saves the state of a sampled task i.e. an episode
from_dict	Create a configuration instance from a dictionary.
load_config	Loads a configuration from a file
load_from_json	Load the configuration from a JSON file.
model_post_init	Initialize and validate configuration after Pydantic model initialization
save_config	Saves the current configuration to the output directory
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
CameraConfig	type
PolicyConfig	type
RobotConfig	type
benchmark_path	Path | None
camera_config	FrankaDroidCameraSystem
collision_free_pose_limit	int
config_version	str
ctrl_dt_ms	float
data_split	str
datagen_profiler	bool
end_on_success	bool
environment_light_intensity	float
eval_runtime_params	Any
filter_for_successful_trajectories	bool
fps	float
log_level	str
num_envs	int
num_workers	int
output_dir	Path
policy_config	PickPlannerPolicyConfig
policy_dt_ms	float
profile	bool
profiler	Profiler | None
robot_config	FloatingRUMRobotConfig
scene_dataset	str
seed	int | None
sim_dt_ms	float
tag	str	A string describing the experiment.
task_config	PickTaskConfig
task_config_preset	PickTaskConfig | None
task_config_preset_exp	AllTaskConfigs | None
task_config_preset_scn	AllTaskConfigs | None
task_horizon	int | None
task_sampler_config	RUMPickTaskSamplerConfig
task_type	str
use_passive_viewer	bool
use_wandb	bool
viewer_cam_dict	dict
viewer_camera	None
wandb_name	str | None
wandb_project	str | None

CameraConfig class-attribute

CameraConfig: type = CameraSystemConfig

PolicyConfig class-attribute

PolicyConfig: type = BasePolicyConfig

RobotConfig class-attribute

RobotConfig: type = BaseRobotConfig

benchmark_path class-attribute instance-attribute

benchmark_path: Path | None = None

camera_config class-attribute instance-attribute

camera_config: FrankaDroidCameraSystem = FrankaRandomizedD405D455CameraSystem(img_resolution=(960, 720))

collision_free_pose_limit class-attribute instance-attribute

collision_free_pose_limit: int = 3

config_version class-attribute instance-attribute

config_version: str = '0.1'

ctrl_dt_ms class-attribute instance-attribute

ctrl_dt_ms: float = 2.0

data_split class-attribute instance-attribute

data_split: str = 'train'

datagen_profiler class-attribute instance-attribute

datagen_profiler: bool = True

end_on_success class-attribute instance-attribute

end_on_success: bool = False

environment_light_intensity class-attribute instance-attribute

environment_light_intensity: float = 15000.0

eval_runtime_params class-attribute instance-attribute

eval_runtime_params: Any = None

filter_for_successful_trajectories class-attribute instance-attribute

filter_for_successful_trajectories: bool = True

fps property

fps: float

log_level class-attribute instance-attribute

log_level: str = 'info'

num_envs class-attribute instance-attribute

num_envs: int = 1

num_workers class-attribute instance-attribute

num_workers: int = 1

output_dir class-attribute instance-attribute

output_dir: Path = ASSETS_DIR / 'experiment_output' / 'datagen' / 'rum_pick_v1'

policy_config class-attribute instance-attribute

policy_config: PickPlannerPolicyConfig = PickPlannerPolicyConfig()

policy_dt_ms class-attribute instance-attribute

policy_dt_ms: float = 66.0

profile class-attribute instance-attribute

profile: bool = True

profiler class-attribute instance-attribute

profiler: Profiler | None = None

robot_config class-attribute instance-attribute

robot_config: FloatingRUMRobotConfig = FloatingRUMRobotConfig()

scene_dataset class-attribute instance-attribute

scene_dataset: str = 'holodeck-objaverse'

seed class-attribute instance-attribute

seed: int | None = None

sim_dt_ms class-attribute instance-attribute

sim_dt_ms: float = 2.0

tag property

tag: str

A string describing the experiment.

task_config class-attribute instance-attribute

task_config: PickTaskConfig = PickTaskConfig(task_cls=PickTask)

task_config_preset class-attribute instance-attribute

task_config_preset: PickTaskConfig | None = None

task_config_preset_exp class-attribute instance-attribute

task_config_preset_exp: AllTaskConfigs | None = None

task_config_preset_scn class-attribute instance-attribute

task_config_preset_scn: AllTaskConfigs | None = None

task_horizon class-attribute instance-attribute

task_horizon: int | None = 500

task_sampler_config class-attribute instance-attribute

task_sampler_config: RUMPickTaskSamplerConfig = RUMPickTaskSamplerConfig(task_sampler_class=PickTaskSampler, robot_object_z_offset=0)

task_type class-attribute instance-attribute

task_type: str = 'pick'

use_passive_viewer class-attribute instance-attribute

use_passive_viewer: bool = False

use_wandb class-attribute instance-attribute

use_wandb: bool = False

viewer_cam_dict class-attribute instance-attribute

viewer_cam_dict: dict = {'distance': 5.0, 'azimuth': 45.0, 'elevation': -30.0, 'lookat': [0.0, 0.0, 0.5]}

viewer_camera class-attribute instance-attribute

viewer_camera: None = None

wandb_name class-attribute instance-attribute

wandb_name: str | None = None

wandb_project class-attribute instance-attribute

wandb_project: str | None = 'molmo-spaces-data-generation'

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

SavedEpisode

Bases: Config

Config informationd describing a sinlge episode

Classes:

Name	Description
Config

Methods:

Name	Description
from_dict	Create a configuration instance from a dictionary.
load_from_json	Load the configuration from a JSON file.
save_to_json	Save the configuration to a JSON file.
to_dict	Convert the configuration to a dictionary.
to_json

Attributes:

Name	Type	Description
camera_config	AllCameraSystems | None
robot_config	FrankaRobotConfig | None
task_cls_str	str | None
task_config	PickTaskConfig | None

camera_config class-attribute instance-attribute

camera_config: AllCameraSystems | None = None

robot_config class-attribute instance-attribute

robot_config: FrankaRobotConfig | None = None

task_cls_str class-attribute instance-attribute

task_cls_str: str | None = None

task_config class-attribute instance-attribute

task_config: PickTaskConfig | None = None

Config

Attributes:

Name	Type	Description
arbitrary_types_allowed

arbitrary_types_allowed class-attribute instance-attribute

arbitrary_types_allowed = True

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

freeze_task_config

freeze_task_config(observation, task: BaseMujocoTask = None) -> None

Saves the state of a sampled task i.e. an episode

Source code in molmo_spaces/configs/abstract_exp_config.py

def freeze_task_config(self, observation, task: BaseMujocoTask = None) -> None:
    """Saves the state of a sampled task i.e. an episode"""
    sc = self.SavedEpisode()

    # RMH: deep argument VERY IMPORTANT. Mutates config for future episodes otherwise
    sc.robot_config = self.robot_config.model_copy(deep=True)
    # remove un-serializable
    sc.robot_config.robot_cls = None
    sc.robot_config.robot_factory = None
    sc.robot_config.robot_view_factory = None
    # save state
    sc.robot_config.init_qpos_noise_range = None  # remove ranges
    sc.robot_config.init_qpos = observation[0]["qpos"]
    sc.camera_config = self.camera_config.model_copy(deep=True)
    for i, camera in enumerate(sc.camera_config.cameras):
        # Some cameras can contain random sampling, e.g. of positions
        # Read the camera's positions and convert them to fixed cameras
        if isinstance(camera, MjcfCameraConfig | RobotMountedCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = RobotMountedCameraConfig(
                name=cam.name,
                reference_body_names=list(cam.reference_body_names),
                camera_offset=list(cam.camera_offset),
                lookat_offset=list(cam.lookat_offset),
                camera_quaternion=list(cam.camera_quaternion),
                fov=cam.fov,
            )
            sc.camera_config.cameras[i] = new_camera

        elif isinstance(camera, RandomizedExocentricCameraConfig | FixedExocentricCameraConfig):
            cam = task.env.camera_manager.registry[camera.name]
            new_camera = FixedExocentricCameraConfig(
                name=cam.name,
                fov=cam.fov,
                pos=list(cam.pos),
                up=list(cam.up),
                forward=list(cam.forward),
            )
            sc.camera_config.cameras[i] = new_camera
        else:
            raise NotImplementedError(f"Cannot freeze camera of type {type(camera).__name__}")

    # for all task relevant objects, save the poses
    # assert task.config.task_config.object_poses is None
    obj_poses = {}
    om = task.env.object_managers[task.env.current_batch_index]
    task_objects = om.get_mobile_objects()
    for task_object in task_objects:
        obj_poses[task_object.name] = pose_mat_to_7d(task_object.pose).tolist()
    task.config.task_config.object_poses = obj_poses

    sc.task_config = self.task_config.model_copy(deep=True)
    # remove un-serializable
    sc.task_config.task_cls = None
    # save the name of the task class
    sc.task_cls_str = (
        self.task_config.task_cls.__module__ + "." + self.task_config.task_cls.__name__
    )

    assert sc.task_config.robot_base_pose is not None

    sc_bytes = pickle.dumps(sc)
    sc_b64 = base64.b64encode(sc_bytes).decode("utf-8")
    return sc_b64

from_dict classmethod

from_dict(data: dict) -> Config

Create a configuration instance from a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def from_dict(cls, data: dict) -> "Config":
    """Create a configuration instance from a dictionary."""
    return cls.model_validate(data)

load_config staticmethod

load_config(output_dir: Path) -> MlSpacesExpConfig

Loads a configuration from a file

Source code in molmo_spaces/configs/abstract_exp_config.py

@staticmethod
def load_config(output_dir: Path) -> MlSpacesExpConfig:
    """Loads a configuration from a file"""
    config_path = output_dir / "experiment_config.pkl"
    with open(config_path, "rb") as f:
        config = pickle.load(f)
    log.info(f"Loaded experiment configuration from {output_dir}")
    return config

load_from_json classmethod

load_from_json(file_path: str) -> Config

Load the configuration from a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

@classmethod
def load_from_json(cls, file_path: str) -> "Config":
    """Load the configuration from a JSON file."""
    with open(file_path, "r") as f:
        data = f.read()
    return cls.model_validate_json(data)

model_post_init

model_post_init(__context) -> None

Initialize and validate configuration after Pydantic model initialization

Source code in molmo_spaces/configs/base_pick_config.py

def model_post_init(self, __context) -> None:
    """Initialize and validate configuration after Pydantic model initialization"""
    super().model_post_init(__context)

    try:
        self.policy_config = self._init_policy_config()
    except RuntimeError as e:
        # Check if this is a CUDA/GPU-related error
        error_msg = str(e)
        if "NVIDIA" in error_msg or "CUDA" in error_msg or "GPU" in error_msg:
            # No GPU available - this is expected on manager nodes that just coordinate jobs
            # Policy config will be initialized later on worker nodes that have GPUs
            print(
                f"Warning: Skipping policy config initialization due to missing GPU: {error_msg}"
            )
            self.policy_config = None
        else:
            raise

    # Auto-create profiler instance if profiling is enabled
    if self.profile and self.profiler is None:
        self.profiler = Profiler()

save_config

save_config(output_dir=None) -> None

Saves the current configuration to the output directory

Source code in molmo_spaces/configs/abstract_exp_config.py

def save_config(self, output_dir=None) -> None:
    """Saves the current configuration to the output directory"""
    if output_dir is None:
        output_dir = self.output_dir
    output_dir = Path(output_dir)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"experiment_config_{timestamp}.pkl"
    with open(config_path, "wb") as f:
        pickle.dump(self, f)
    log.info(f"Saved experiment configuration to {output_dir}")

save_to_json

save_to_json(file_path: str) -> None

Save the configuration to a JSON file.

Source code in molmo_spaces/configs/abstract_config.py

def save_to_json(self, file_path: str) -> None:
    """Save the configuration to a JSON file."""
    with open(file_path, "w") as f:
        f.write(self.to_json())

to_dict

to_dict() -> dict

Convert the configuration to a dictionary.

Source code in molmo_spaces/configs/abstract_config.py

def to_dict(self) -> dict:
    """Convert the configuration to a dictionary."""
    return self.model_dump()

to_json

to_json() -> str

Source code in molmo_spaces/configs/abstract_config.py

def to_json(self) -> str:
    # TODO(max): this can cause errors, try printing out the config to see missmatches w/ print(self)
    return self.model_dump_json(warnings="error")

config_registry

Configuration registry for data generation experiments.

This registry allows multiple config classes to be defined in any file structure, without strict naming conventions. Simply register your config classes here.

Functions:

Name	Description
get_config_class	Get a config class by name from the registry.
get_registry_size	Get the number of registered configs.
list_available_configs	List all available config names in the registry.
register_config	Decorator to register a config class in the registry.

Attributes:

Name	Type	Description
log

log module-attribute

log = getLogger(__name__)

get_config_class

get_config_class(name: str) -> type[MlSpacesExpConfig]

Get a config class by name from the registry.

Parameters:

Name	Type	Description	Default
name	str	The name of the config class	required

Returns:

Type	Description
type[MlSpacesExpConfig]	The config class

Raises:

Type	Description
ValueError	If the config name is not found

Source code in molmo_spaces/data_generation/config_registry.py

def get_config_class(name: str) -> type[MlSpacesExpConfig]:
    """Get a config class by name from the registry.

    Args:
        name: The name of the config class

    Returns:
        The config class

    Raises:
        ValueError: If the config name is not found
    """
    if name not in _MJT_CONFIG_REGISTRY:
        available = list(_MJT_CONFIG_REGISTRY.keys())
        raise ValueError(f"Config '{name}' not found. Available configs: {available}")
    return _MJT_CONFIG_REGISTRY[name]

get_registry_size

get_registry_size() -> int

Get the number of registered configs.

Source code in molmo_spaces/data_generation/config_registry.py

def get_registry_size() -> int:
    """Get the number of registered configs."""
    return len(_MJT_CONFIG_REGISTRY)

list_available_configs

list_available_configs() -> list[str]

List all available config names in the registry.

Source code in molmo_spaces/data_generation/config_registry.py

def list_available_configs() -> list[str]:
    """List all available config names in the registry."""
    return list(_MJT_CONFIG_REGISTRY.keys())

register_config

register_config(name: str, strict: bool = True)

Decorator to register a config class in the registry.

Usage

@register_config("MyExperimentConfig") class MyExperimentConfig(MlSpacesExpConfig): ...

Parameters:

Name	Type	Description	Default
name	str	The name to register this config under (used in command line)	required

Source code in molmo_spaces/data_generation/config_registry.py

def register_config(name: str, strict: bool = True):
    """Decorator to register a config class in the registry.

    Usage:
        @register_config("MyExperimentConfig")
        class MyExperimentConfig(MlSpacesExpConfig):
            ...

    Args:
        name: The name to register this config under (used in command line)
    """

    def decorator(cls: type[MlSpacesExpConfig]):
        if name in _MJT_CONFIG_REGISTRY:
            existing_cls = _MJT_CONFIG_REGISTRY[name]
            existing_cls_id = f"{existing_cls.__module__}.{existing_cls.__name__}"
            new_cls_id = f"{cls.__module__}.{cls.__name__}"
            if strict:
                raise ValueError(
                    f"Config '{name}' already registered as {existing_cls_id}, trying to register as {new_cls_id}"
                )
            log.warning(
                f"Overriding existing config '{name}'. Was {existing_cls_id}, now {new_cls_id}"
            )
        _MJT_CONFIG_REGISTRY[name] = cls
        return cls

    return decorator

main

Functions:

Name	Description
auto_import_configs	Auto-import all config files so they register themselves
get_args
main

auto_import_configs

auto_import_configs() -> None

Auto-import all config files so they register themselves

Source code in molmo_spaces/data_generation/main.py

def auto_import_configs() -> None:
    """Auto-import all config files so they register themselves"""
    # Get the config directory path
    current_dir = os.path.dirname(__file__)
    config_dir = os.path.join(current_dir, "config")

    if not os.path.exists(config_dir):
        print(f"Warning: Config directory not found: {config_dir}")
        return

    # Import all .py files in the config directory
    config_files = glob.glob(os.path.join(config_dir, "*.py"))

    for config_path in config_files:
        # Skip __init__.py
        if config_path.endswith("__init__.py"):
            continue

        # Load the module with the full module path for proper pickling
        module_filename = os.path.splitext(os.path.basename(config_path))[0]
        full_module_name = f"molmo_spaces.data_generation.config.{module_filename}"

        # Use standard import instead of spec_from_file_location
        # This ensures the module has the correct __name__ for pickling
        try:
            importlib.import_module(full_module_name)
        except Exception as e:
            print(f"Warning: Could not load config from {full_module_name}: {e}")
            continue

get_args

get_args()

Source code in molmo_spaces/data_generation/main.py

def get_args():
    parser = argparse.ArgumentParser(
        description="MolmoSpaces data generation pipeline",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "exp_config_cls",
        type=str,
        help="Name of the experiment config class to use (e.g., FrankaPickDroidDataGenConfig), "
        "optionally with the module name prepended with a colon (e.g. molmo_spaces.data_generation.config.object_manipulation_datagen_configs:FrankaPickDroidDataGenConfig). "
        "If the module is specified, only that module will be imported to populate the registry. Otherwise, all config files will be imported.",
    )
    return parser.parse_args()

main

main() -> None

Source code in molmo_spaces/data_generation/main.py

def main() -> None:
    args = get_args()
    exp_config_cls = args.exp_config_cls

    # np.random.seed(42)

    if (
        ":" in exp_config_cls
    ):  # if the module is specified, import it and use the class from that module
        exp_config_module, exp_config_cls = exp_config_cls.split(":")
        importlib.import_module(exp_config_module)
    else:  # otherwise, auto-import all config files to populate the registry
        auto_import_configs()

    # Get the config class from the registry
    ExpConfigClass = get_config_class(exp_config_cls)
    # NOTE: You may optionally load additional exp_config arguments from argparse command line
    exp_config_args = vars(args)
    # pop exp_config_cls from args
    exp_config_args.pop("exp_config_cls")
    # Create exp_config instance with args
    exp_config = ExpConfigClass(**exp_config_args)

    # Optional: Modify the config parameters here if needed
    # Eg. for hyperparamter sweeps etc.

    # Generate unique run name
    exp_config_name = exp_config_cls  # Use the class name directly

    # Determine output directory structure
    # For local debugging (non-shared paths), add timestamp to avoid collisions
    # For production (datagen output targets shared filesystem paths), use simple structure
    is_shared_fs_path = "/mnt/shared" in str(
        exp_config.output_dir
    )  # or whatever your mount point is

    if is_shared_fs_path:
        # Production: simple structure without timestamp
        exp_config.output_dir = exp_config.output_dir / exp_config_name
    else:
        # Local debug: add timestamp to avoid collisions
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        exp_config.output_dir = exp_config.output_dir / exp_config_name / timestamp

    os.makedirs(exp_config.output_dir, exist_ok=True)

    # Initialize wandb if enabled - with auto run name
    if exp_config.use_wandb:
        import wandb

        if exp_config.wandb_name is None:
            # Generate timestamp for wandb run name
            wandb_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            exp_config.wandb_name = f"{exp_config_name}_{wandb_timestamp}"
        wandb.init(
            project=exp_config.wandb_project, name=exp_config.wandb_name, config=vars(exp_config)
        )

    exp_config.save_config()

    # Create rollout runner with the set config parameters
    runner = ParallelRolloutRunner(exp_config)

    success_count, total_count = runner.run()
    print(f"Success count: {success_count}, Total count: {total_count}")

    # Close wandb run
    if exp_config.use_wandb:
        wandb.finish()

pipeline

Classes:

Name	Description
ParallelRolloutRunner	Orchestrates parallel house processing for offline data generation using multiprocessing.

Functions:

Name	Description
cleanup_context	Context manager to ensure proper cleanup of MuJoCo resources
cleanup_episode_resources	Cleanup resources after an episode.
get_detailed_memory_info	Get detailed memory usage information
get_process_memory	Get current memory usage of the process in MB
house_processing_worker	Standalone worker function that processes work items sequentially from a shared counter.
log_detailed_memory	Log detailed memory usage information
log_memory_usage	Log current memory usage
save_house_trajectories	Batch and save trajectory data for a house.
setup_house_dirs	Setup output directories and check for existing output.
setup_policy	Create or return policy for episode.
setup_viewer	Setup passive viewer if configured.

Attributes:

Name	Type	Description
mp_context

mp_context module-attribute

mp_context = get_context('forkserver') if is_available() else get_context('spawn')

ParallelRolloutRunner

ParallelRolloutRunner(exp_config: MlSpacesExpConfig)

Orchestrates parallel house processing for offline data generation using multiprocessing.

This class is designed for data generation workloads where: - Each worker processes complete houses independently - Task sampling involves heavy Python operations (scene setup, randomization) - Workers can run in separate processes for true parallelism

Note: This is would be an inefficient runner for RL training, and is only intended for data generation or online evaluation tasks. For RL with vectorized environments, consider implementing a separate RLRolloutRunner that works without the multiprocessing overhead, and move threading into env.py with a MJXMuJoCoEnv class or similar.

Customization via Subclassing

You can customize rollout behavior by subclassing this runner and overriding the static methods run_single_rollout and/or process_single_house.

Example: class CustomRolloutRunner(ParallelRolloutRunner): @staticmethod def run_single_rollout(episode_seed, task, policy, kwargs): # Add custom logging or behavior print(f"Starting rollout with seed {episode_seed}") return ParallelRolloutRunner.run_single_rollout( episode_seed, task, policy, kwargs )

# Use your custom runner
runner = CustomRolloutRunner(exp_config)
runner.run()  # Workers will use CustomRolloutRunner.run_single_rollout!

Initialize the parallel rollout runner.

Parameters:

Name	Type	Description	Default
exp_config	MlSpacesExpConfig	Experiment configuration	required

Methods:

Name	Description
get_episode_seed	Get seed for episode.
get_episode_spec_at_index	Get episode specification at given index.
get_episode_task_sampler	Get task sampler for episode. Override to create per-episode samplers.
get_max_episode_attempts	Get maximum number of episode attempts for this house.
load_episodes_for_house	Load episode specifications for a house.
prepare_episode_config	Prepare config for a specific episode. Override to modify config per-episode.
process_single_house	Process all episodes for a single house using customizable hooks.
run	Run house-by-house rollouts using multiprocessing workers.
run_single_rollout	Execute a single rollout with the given task and policy.
sample_task_from_spec	Sample task from specification.
should_close_episode_task_sampler	Whether to close task sampler after each episode.
should_stop_early	Whether to stop before processing all episodes (e.g., enough successes).

Attributes:

Name	Type	Description
completed_houses
config
counter_lock
house_counter
house_indices
logger
max_allowed_sequential_irrecoverable_failures
max_allowed_sequential_rollout_failures
max_allowed_sequential_task_sampler_failures
profiler
samples_per_house
shutdown_event
skipped_houses
success_count
total_count
total_houses
wandb_enabled
work_items	list[tuple[int, int, int, int]]

Source code in molmo_spaces/data_generation/pipeline.py

def __init__(self, exp_config: MlSpacesExpConfig) -> None:
    """
    Initialize the parallel rollout runner.

    Args:
        exp_config: Experiment configuration
    """
    self.config = exp_config

    # House-based processing setup
    self.house_indices = exp_config.task_sampler_config.house_inds
    self.samples_per_house = exp_config.task_sampler_config.samples_per_house

    # don't have houses specified, use all, need to find out which ones those are
    if self.house_indices is None:
        # if user-specified scene xml paths, use all of them
        if exp_config.task_sampler_config.scene_xml_paths:
            assert exp_config.scene_dataset == "user"
            self.house_indices = list(
                range(len(exp_config.task_sampler_config.scene_xml_paths))
            )
        else:
            # For normal datagen: use all houses from scene mapping
            mapping = get_scenes(exp_config.scene_dataset, exp_config.data_split)
            self.house_indices = [
                k for k, v in mapping[exp_config.data_split].items() if v is not None
            ]

    self.total_houses = len(self.house_indices)

    # Build (house_id, batch_samples, batch_num, total_batches) work items
    episodes_per_batch = exp_config.task_sampler_config.episodes_per_batch
    self.work_items: list[tuple[int, int, int, int]] = []
    for house_id in self.house_indices:
        total_batches = max(1, round(self.samples_per_house / episodes_per_batch))
        base = self.samples_per_house // total_batches
        extra = self.samples_per_house % total_batches
        for b in range(total_batches):
            batch_samples = base + (1 if b < extra else 0)
            self.work_items.append((house_id, batch_samples, b + 1, total_batches))

    # Failure tracking limits
    self.max_allowed_sequential_task_sampler_failures = (
        exp_config.task_sampler_config.max_allowed_sequential_task_sampler_failures
    )
    self.max_allowed_sequential_rollout_failures = (
        exp_config.task_sampler_config.max_allowed_sequential_rollout_failures
    )
    self.max_allowed_sequential_irrecoverable_failures = (
        exp_config.task_sampler_config.max_allowed_sequential_irrecoverable_failures
    )

    # Setup shared state for multiprocessing
    self.counter_lock = mp_context.Lock()
    self.shutdown_event = mp_context.Event()
    self.house_counter = mp_context.Value("i", 0)
    self.success_count = mp_context.Value("i", 0)
    self.total_count = mp_context.Value("i", 0)
    self.completed_houses = mp_context.Value("i", 0)
    self.skipped_houses = mp_context.Value("i", 0)

    # SIGTERM handling (only register in main thread)
    # Signal handlers can only be registered in the main thread of the main interpreter
    # This can fail when ParallelRolloutRunner is created from worker threads (e.g., DDP, PyTorch Lightning)
    if threading.current_thread() is threading.main_thread():
        signal.signal(signal.SIGTERM, self._handle_sigterm)
    else:
        # If not in main thread, skip signal handler registration
        # This is safe - the runner will still work, just without SIGTERM handling
        pass

    # Logging & Profiling
    init_logging(
        human_log_level=exp_config.log_level, log_file=exp_config.output_dir / "running_log.log"
    )
    self.logger = get_logger()
    self.profiler = exp_config.profiler

    # WandB initialization (optional, based on environment variables)
    self.wandb_enabled = False
    if "WANDB_RUN_NAME" in os.environ and "WANDB_PROJECT_NAME" in os.environ:
        self.logger.info("Initializing WandB logging...")
        try:
            wandb.init(
                project=os.environ["WANDB_PROJECT_NAME"],
                entity=os.environ.get("WANDB_ENTITY", None),
                name=os.environ["WANDB_RUN_NAME"],
                config={
                    "num_workers": exp_config.num_workers,
                    "total_houses": self.total_houses,
                    "samples_per_house": self.samples_per_house,
                    "total_work_items": len(self.work_items),
                    "total_expected_episodes": sum(wi[1] for wi in self.work_items),
                    "output_dir": str(exp_config.output_dir),
                    "filter_for_successful_trajectories": exp_config.filter_for_successful_trajectories,
                },
            )
            self.wandb_enabled = True
            self.logger.info("WandB initialized successfully")
        except Exception as e:
            self.logger.warning(f"WandB initialization failed: {e}. Continuing without WandB.")
            self.wandb_enabled = False
    else:
        self.logger.info("WandB logging disabled (environment variables not set)")

completed_houses instance-attribute

completed_houses = Value('i', 0)

config instance-attribute

config = exp_config

counter_lock instance-attribute

counter_lock = Lock()

house_counter instance-attribute

house_counter = Value('i', 0)

house_indices instance-attribute

house_indices = house_inds

logger instance-attribute

logger = get_logger()

max_allowed_sequential_irrecoverable_failures instance-attribute

max_allowed_sequential_irrecoverable_failures = max_allowed_sequential_irrecoverable_failures

max_allowed_sequential_rollout_failures instance-attribute

max_allowed_sequential_rollout_failures = max_allowed_sequential_rollout_failures

max_allowed_sequential_task_sampler_failures instance-attribute

max_allowed_sequential_task_sampler_failures = max_allowed_sequential_task_sampler_failures

profiler instance-attribute

profiler = profiler

samples_per_house instance-attribute

samples_per_house = samples_per_house

shutdown_event instance-attribute

shutdown_event = Event()

skipped_houses instance-attribute

skipped_houses = Value('i', 0)

success_count instance-attribute

success_count = Value('i', 0)

total_count instance-attribute

total_count = Value('i', 0)

total_houses instance-attribute

total_houses = len(house_indices)

wandb_enabled instance-attribute

wandb_enabled = True

work_items instance-attribute

work_items: list[tuple[int, int, int, int]] = []

get_episode_seed staticmethod

get_episode_seed(episode_idx: int, episode_spec, task_sampler) -> int

Get seed for episode.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def get_episode_seed(episode_idx: int, episode_spec, task_sampler) -> int:
    """Get seed for episode."""
    return task_sampler.current_seed

get_episode_spec_at_index staticmethod

get_episode_spec_at_index(episode_specs: list, idx: int) -> Any

Get episode specification at given index.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def get_episode_spec_at_index(episode_specs: list, idx: int) -> Any:
    """Get episode specification at given index."""
    return episode_specs[idx]

get_episode_task_sampler staticmethod

get_episode_task_sampler(exp_config: MlSpacesExpConfig, episode_spec, shared_task_sampler, datagen_profiler: DatagenProfiler | None) -> Any

Get task sampler for episode. Override to create per-episode samplers.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def get_episode_task_sampler(
    exp_config: MlSpacesExpConfig,
    episode_spec,
    shared_task_sampler,
    datagen_profiler: DatagenProfiler | None,
) -> Any:
    """Get task sampler for episode. Override to create per-episode samplers."""
    return shared_task_sampler

get_max_episode_attempts staticmethod

get_max_episode_attempts(episode_specs: list, samples_per_house: int, exp_config: MlSpacesExpConfig) -> int

Get maximum number of episode attempts for this house.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def get_max_episode_attempts(
    episode_specs: list,
    samples_per_house: int,
    exp_config: MlSpacesExpConfig,
) -> int:
    """Get maximum number of episode attempts for this house."""
    return len(episode_specs)

load_episodes_for_house staticmethod

load_episodes_for_house(exp_config: MlSpacesExpConfig, house_id: int, batch_suffix: str, worker_task_sampler, worker_logger) -> tuple[list, Any]

Load episode specifications for a house.

Returns:

Name	Type	Description
tuple	tuple[list, Any]	(episode_specs, task_sampler_to_use) - episode_specs: List of saved configs or None values - task_sampler_to_use: Task sampler for sampling tasks

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def load_episodes_for_house(
    exp_config: MlSpacesExpConfig,
    house_id: int,
    batch_suffix: str,
    worker_task_sampler,
    worker_logger,
) -> tuple[list, Any]:
    """
    Load episode specifications for a house.

    Returns:
        tuple: (episode_specs, task_sampler_to_use)
            - episode_specs: List of saved configs or None values
            - task_sampler_to_use: Task sampler for sampling tasks
    """
    # Datagen mode: generate None list for fresh sampling
    max_multiplier = exp_config.task_sampler_config.max_total_attempts_multiplier
    num_samples = exp_config.task_sampler_config.samples_per_house * max_multiplier
    return [None] * num_samples, worker_task_sampler

prepare_episode_config staticmethod

prepare_episode_config(exp_config: MlSpacesExpConfig, episode_spec, episode_idx: int) -> MlSpacesExpConfig

Prepare config for a specific episode. Override to modify config per-episode.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def prepare_episode_config(
    exp_config: MlSpacesExpConfig,
    episode_spec,
    episode_idx: int,
) -> MlSpacesExpConfig:
    """Prepare config for a specific episode. Override to modify config per-episode."""
    return exp_config  # Base class doesn't modify config

process_single_house staticmethod

process_single_house(worker_id: int, worker_logger, house_id: int, exp_config: MlSpacesExpConfig, samples_per_house: int, shutdown_event, task_sampler, preloaded_policy: BasePolicy | None = None, max_allowed_sequential_task_sampler_failures: int = 10, max_allowed_sequential_rollout_failures: int = 10, filter_for_successful_trajectories: bool = False, runner_class=None, batch_num: int | None = None, total_batches: int | None = None, datagen_profiler: DatagenProfiler | None = None) -> tuple[int, int, bool]

Process all episodes for a single house using customizable hooks.

This method uses a while loop to iterate over episodes, calling hook methods via runner_class to allow subclasses to customize behavior without duplicating the entire method.

Hooks called (override in subclass to customize): - load_episodes_for_house: Load episode specs from source (JSON, etc.) - get_max_episode_attempts: Maximum iterations of the episode loop - should_stop_early: Whether to stop before max attempts (e.g., enough successes) - prepare_episode_config: Modify config per-episode - get_episode_task_sampler: Get/create task sampler for episode - sample_task_from_spec: Sample task from specification - get_episode_seed: Get seed for episode - should_close_episode_task_sampler: Whether to close sampler per-episode

Parameters:

Name	Type	Description	Default
worker_id	int	ID of the worker thread/process	required
worker_logger		Logger instance for this worker	required
house_id	int	Index of the house to process	required
exp_config	MlSpacesExpConfig	Experiment configuration	required
samples_per_house	int	Number of episodes to collect for this house	required
shutdown_event		Event to signal shutdown	required
task_sampler		Task sampler instance (shared across houses for this worker)	required
preloaded_policy	BasePolicy | None	Optional pre-initialized policy instance	None
max_allowed_sequential_task_sampler_failures	int	Max consecutive task sampling failures	10
max_allowed_sequential_rollout_failures	int	Max consecutive rollout failures	10
filter_for_successful_trajectories	bool	Whether to filter for successful trajectories only	False
runner_class		Runner class with hook methods to call	None
batch_num	int | None	Batch number for this house (for batched processing)	None
total_batches	int | None	Total number of batches for this house	None
datagen_profiler	DatagenProfiler | None	DatagenProfiler for per-worker timing (optional)	None

Returns:

Name	Type	Description
tuple	tuple[int, int, bool]	(house_success_count, house_total_count, irrecoverable_failure_flag)

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def process_single_house(
    worker_id: int,
    worker_logger,
    house_id: int,
    exp_config: MlSpacesExpConfig,
    samples_per_house: int,
    shutdown_event,
    task_sampler,
    preloaded_policy: BasePolicy | None = None,
    max_allowed_sequential_task_sampler_failures: int = 10,
    max_allowed_sequential_rollout_failures: int = 10,
    filter_for_successful_trajectories: bool = False,
    runner_class=None,
    batch_num: int | None = None,
    total_batches: int | None = None,
    datagen_profiler: DatagenProfiler | None = None,
) -> tuple[int, int, bool]:
    """
    Process all episodes for a single house using customizable hooks.

    This method uses a while loop to iterate over episodes, calling hook methods
    via runner_class to allow subclasses to customize behavior without duplicating
    the entire method.

    Hooks called (override in subclass to customize):
    - load_episodes_for_house: Load episode specs from source (JSON, etc.)
    - get_max_episode_attempts: Maximum iterations of the episode loop
    - should_stop_early: Whether to stop before max attempts (e.g., enough successes)
    - prepare_episode_config: Modify config per-episode
    - get_episode_task_sampler: Get/create task sampler for episode
    - sample_task_from_spec: Sample task from specification
    - get_episode_seed: Get seed for episode
    - should_close_episode_task_sampler: Whether to close sampler per-episode

    Args:
        worker_id: ID of the worker thread/process
        worker_logger: Logger instance for this worker
        house_id: Index of the house to process
        exp_config: Experiment configuration
        samples_per_house: Number of episodes to collect for this house
        shutdown_event: Event to signal shutdown
        task_sampler: Task sampler instance (shared across houses for this worker)
        preloaded_policy: Optional pre-initialized policy instance
        max_allowed_sequential_task_sampler_failures: Max consecutive task sampling failures
        max_allowed_sequential_rollout_failures: Max consecutive rollout failures
        filter_for_successful_trajectories: Whether to filter for successful trajectories only
        runner_class: Runner class with hook methods to call
        batch_num: Batch number for this house (for batched processing)
        total_batches: Total number of batches for this house
        datagen_profiler: DatagenProfiler for per-worker timing (optional)

    Returns:
        tuple: (house_success_count, house_total_count, irrecoverable_failure_flag)
    """
    house_success_count = 0
    house_total_count = 0
    irrecoverable_failure_in_house = False

    # Setup directories and check for existing output
    house_output_dir, house_debug_dir, batch_suffix, should_skip = setup_house_dirs(
        exp_config, house_id, batch_num, total_batches
    )
    if should_skip:
        worker_logger.info(
            f"SKIPPING HOUSE {house_id} BATCH {batch_num}/{total_batches}: "
            f"Output already exists at {house_output_dir / f'trajectories{batch_suffix}.h5'}"
        )
        return 0, 0, False

    # Load episodes using hook - allows subclasses to load from different scene datasets
    episode_specs, shared_task_sampler = runner_class.load_episodes_for_house(
        exp_config, house_id, batch_suffix, task_sampler, worker_logger
    )

    if not episode_specs:
        worker_logger.warning(f"No episodes to process for house {house_id}")
        return 0, 0, False

    max_attempts = runner_class.get_max_episode_attempts(
        episode_specs, samples_per_house, exp_config
    )

    # Collect raw history data for this house
    house_raw_histories = []
    house_debug_raw_histories = []

    # Sequential failure tracking
    num_sequential_task_sampler_failures = 0
    num_sequential_rollout_failures = 0
    viewer = None

    # While loop with explicit index - allows subclasses to customize iteration
    episode_idx = 0
    while episode_idx < max_attempts:
        # Check early stop condition (e.g., enough successes for datagen)
        should_stop = runner_class.should_stop_early(
            len(house_raw_histories), samples_per_house, exp_config=exp_config
        )
        if should_stop:
            break

        # Check for shutdown signal
        if shutdown_event.is_set():
            worker_logger.info(f"Worker {worker_id} house {house_id} received shutdown signal")
            irrecoverable_failure_in_house = True
            break

        # Check for too many consecutive task sampling failures
        if num_sequential_task_sampler_failures >= max_allowed_sequential_task_sampler_failures:
            worker_logger.error(
                f"Worker {worker_id} house {house_id} encountered "
                f"{num_sequential_task_sampler_failures} consecutive task sampling failures. "
                "This is unrecoverable."
            )
            irrecoverable_failure_in_house = True
            break

        # Check for too many consecutive rollout failures
        if num_sequential_rollout_failures >= max_allowed_sequential_rollout_failures:
            worker_logger.error(
                f"Worker {worker_id} house {house_id} rollout failed across "
                f"{num_sequential_rollout_failures} retries. This is irrecoverable."
            )
            irrecoverable_failure_in_house = True
            break

        # Get episode spec for this iteration
        episode_spec = runner_class.get_episode_spec_at_index(episode_specs, episode_idx)

        # Track state for this episode
        task = None
        policy = None
        episode_task_sampler = None
        success = False
        task_sampling_failed = False
        house_invalid = False

        if datagen_profiler is not None:
            datagen_profiler.start("episode_total")

        # Prepare episode-specific config
        episode_config = runner_class.prepare_episode_config(
            exp_config, episode_spec, episode_idx
        )

        with cleanup_context():
            if viewer is not None:
                viewer.close()
                viewer = None

            # Task sampling phase
            task_sampling_start = time.perf_counter()

            try:
                # Get task sampler for this episode (shared or per-episode)
                episode_task_sampler = runner_class.get_episode_task_sampler(
                    episode_config, episode_spec, shared_task_sampler, datagen_profiler
                )

                # Sample task
                task = runner_class.sample_task_from_spec(
                    episode_task_sampler, house_id, episode_spec, episode_idx
                )

                if task is None:
                    worker_logger.info(
                        f"Worker {worker_id} house {house_id} episode {episode_idx}: "
                        "task sampling returned None"
                    )
                    house_invalid = True
                else:
                    # Record successful sampling time
                    if datagen_profiler is not None:
                        datagen_profiler.record(
                            "task_sampling", time.perf_counter() - task_sampling_start
                        )
                        task.set_datagen_profiler(datagen_profiler)

                    num_sequential_task_sampler_failures = 0

                    worker_logger.info(
                        f"Worker {worker_id} house {house_id} episode {episode_idx}/{max_attempts} "
                        f"collected={len(house_raw_histories)}/{samples_per_house}"
                    )

            except HouseInvalidForTask as e:
                traceback.print_exc()
                worker_logger.warning(
                    f"Worker {worker_id} house {house_id} episode {episode_idx} "
                    f"HouseInvalidForTask: {e.reason}"
                )
                house_invalid = True
                if datagen_profiler is not None:
                    datagen_profiler.record(
                        "task_sampling_failed", time.perf_counter() - task_sampling_start
                    )

            except Exception as e:
                traceback.print_exc()
                worker_logger.error(
                    f"Worker {worker_id} house {house_id} episode {episode_idx} "
                    f"task sampling error: {str(e)}"
                )
                num_sequential_task_sampler_failures += 1
                task_sampling_failed = True
                if datagen_profiler is not None:
                    datagen_profiler.record(
                        "task_sampling_failed", time.perf_counter() - task_sampling_start
                    )

            # Rollout phase (only if task sampling succeeded)
            if task is not None and not house_invalid and not task_sampling_failed:
                try:
                    # Setup policy and viewer
                    policy = setup_policy(
                        episode_config, task, preloaded_policy, datagen_profiler
                    )
                    viewer = setup_viewer(episode_config, task, policy, viewer)

                    # Get episode seed
                    episode_seed = runner_class.get_episode_seed(
                        episode_idx, episode_spec, episode_task_sampler
                    )

                    # Run the rollout
                    success = runner_class.run_single_rollout(
                        episode_seed=episode_seed,
                        task=task,
                        policy=policy,
                        profiler=episode_config.profiler,
                        viewer=viewer,
                        shutdown_event=shutdown_event,
                        datagen_profiler=datagen_profiler,
                        end_on_success=exp_config.end_on_success,
                    )

                    num_sequential_rollout_failures = 0

                    # Extract object name for logging if available
                    object_name = "unknown"
                    if hasattr(task, "config") and hasattr(task.config, "task_config"):
                        if hasattr(task.config.task_config, "pickup_obj_name"):
                            object_name = task.config.task_config.pickup_obj_name

                    worker_logger.info(
                        f"Worker {worker_id} house {house_id} episode {episode_idx} "
                        f"object {object_name} completed with success={success}"
                    )

                    # Collect trajectory
                    should_save = success or not filter_for_successful_trajectories
                    history = task.get_history()
                    should_save_debug = not should_save and random.random() < 0.01

                    if should_save or should_save_debug:
                        episode_info = {
                            "history": history,
                            "sensor_suite": task.sensor_suite,
                            "success": success,
                            "seed": episode_seed,
                        }
                        if should_save:
                            house_raw_histories.append(episode_info)
                        elif should_save_debug:
                            house_debug_raw_histories.append(episode_info)
                            worker_logger.info(
                                f"Queueing failed trajectory for debug (seed: {episode_seed})"
                            )
                    else:
                        del history

                    # Update house counters
                    house_total_count += 1
                    if success:
                        house_success_count += 1
                    else:
                        # Report failure for this asset (may lead to dynamic blacklisting)
                        asset_uid = task_sampler.get_asset_uid_from_object(
                            task.env, object_name
                        )
                        if asset_uid:
                            task_sampler.report_asset_failure(
                                asset_uid, "rollout failed (e.g., IK failure)"
                            )

                    if datagen_profiler is not None:
                        datagen_profiler.end("episode_total")
                        datagen_profiler.log_episode_summary(
                            episode_idx=episode_idx,
                            house_id=house_id,
                            success=success,
                        )

                except Exception as e:
                    worker_logger.error(
                        f"Worker {worker_id} house {house_id} episode {episode_idx} rollout error: {str(e)}"
                    )
                    traceback.print_exc()
                    num_sequential_rollout_failures += 1

                    # Report failure for this asset (may lead to dynamic blacklisting)
                    try:
                        asset_uid = task_sampler.get_asset_uid_from_object(
                            task.env, object_name
                        )
                        if asset_uid:
                            task_sampler.report_asset_failure(
                                asset_uid, f"rollout exception: {e}"
                            )
                    except Exception:
                        pass  # Don't let failure tracking break the error handling

                    if datagen_profiler is not None:
                        datagen_profiler.end("episode_total")

            else:
                # Task sampling failed or house invalid
                if datagen_profiler is not None:
                    datagen_profiler.end("episode_total")

            # Cleanup resources
            cleanup_episode_resources(
                task=task,
                policy=policy,
                task_sampler=episode_task_sampler,
                preloaded_policy=preloaded_policy,
                close_task_sampler=runner_class.should_close_episode_task_sampler(),
            )

        # Handle house invalid - break after cleanup
        if house_invalid:
            irrecoverable_failure_in_house = True
            break

        # Always increment episode index
        episode_idx += 1

    # Cleanup viewer
    if viewer is not None:
        viewer.close()
        viewer = None

    # Check shutdown signal before saving
    if shutdown_event.is_set():
        worker_logger.info(
            f"Worker {worker_id} house {house_id} shutdown requested, skipping save"
        )
        return house_success_count, house_total_count, True

    # Save trajectories
    save_house_trajectories(
        worker_logger,
        house_raw_histories,
        house_output_dir,
        exp_config,
        batch_suffix,
        datagen_profiler,
        batch_num,
        total_batches,
    )

    # Save debug trajectories
    save_house_trajectories(
        worker_logger,
        house_debug_raw_histories,
        house_debug_dir,
        exp_config,
        batch_suffix,
        datagen_profiler=None,
        batch_num=batch_num,
        total_batches=total_batches,
    )

    worker_logger.info(
        f"Worker {worker_id} completed house {house_id}: "
        f"{house_success_count}/{house_total_count} successful episodes"
    )

    if datagen_profiler is not None:
        datagen_profiler.log_house_summary(
            house_id=house_id,
            success_count=house_success_count,
            total_count=house_total_count,
        )

    return house_success_count, house_total_count, irrecoverable_failure_in_house

run

run(preloaded_policy: BasePolicy | None = None) -> tuple[int, int]

Run house-by-house rollouts using multiprocessing workers.

Parameters:

Name	Type	Description	Default
preloaded_policy	BasePolicy | None	Optional pre-initialized policy instance to use for rollouts. If None, a new policy will be created for each rollout.	None

Returns:

Name	Type	Description
tuple	tuple[int, int]	(success_count, total_count)

Source code in molmo_spaces/data_generation/pipeline.py

def run(self, preloaded_policy: BasePolicy | None = None) -> tuple[int, int]:
    """
    Run house-by-house rollouts using multiprocessing workers.

    Args:
        preloaded_policy: Optional pre-initialized policy instance to use for rollouts.
            If None, a new policy will be created for each rollout.

    Returns:
        tuple: (success_count, total_count)
    """
    total_expected_episodes = sum(wi[1] for wi in self.work_items)
    self.logger.info(
        f"Starting rollout of {self.total_houses} houses "
        f"split into {len(self.work_items)} work items ({total_expected_episodes} total episodes) "
        f"using {self.config.num_workers} worker processes"
    )

    # make a copy of the config in the output directory
    self.logger.info("Evaluation configuration:")
    self.logger.info(pprint.pformat(self.config.model_dump()))
    self.config.save_config(output_dir=Path(self.config.output_dir))

    # Start timing for WandB metrics
    start_time = time.time()

    # Launch worker processes
    if self.config.num_workers > 1:
        processes = []
        for worker_id in range(self.config.num_workers):
            p = mp_context.Process(
                target=house_processing_worker,
                args=(
                    worker_id,
                    self.config,
                    self.work_items,
                    self.shutdown_event,
                    self.counter_lock,
                    self.house_counter,
                    self.success_count,
                    self.total_count,
                    self.completed_houses,
                    self.skipped_houses,
                    self.max_allowed_sequential_task_sampler_failures,
                    self.max_allowed_sequential_rollout_failures,
                    self.max_allowed_sequential_irrecoverable_failures,
                    preloaded_policy,
                    self.config.filter_for_successful_trajectories,
                    type(self),  # Pass the runner class to enable customization via subclassing
                ),
            )
            p.start()
            processes.append(p)

        # Periodic logging loop that monitors progress while workers run
        last_log_time = start_time
        log_interval = 60  # Log every 60 seconds

        while any(p.is_alive() for p in processes):
            # Check if it's time to log
            current_time = time.time()
            if self.wandb_enabled and (current_time - last_log_time) >= log_interval:
                try:
                    # Read current progress from shared counters
                    elapsed_time = current_time - start_time
                    completed = self.completed_houses.value
                    skipped = self.skipped_houses.value
                    success = self.success_count.value
                    total = self.total_count.value
                    active = sum(1 for p in processes if p.is_alive())

                    # Calculate metrics
                    total_work_items = len(self.work_items)
                    success_rate = success / total if total > 0 else 0.0
                    episodes_per_second = total / elapsed_time if elapsed_time > 0 else 0.0
                    completion_percentage = (completed + skipped) / total_work_items * 100

                    # Log to WandB
                    wandb.log(
                        {
                            "elapsed_time_seconds": elapsed_time,
                            "elapsed_time_hours": elapsed_time / 3600,
                            "completed_houses": completed,
                            "skipped_houses": skipped,
                            "success_count": success,
                            "total_count": total,
                            "success_rate": success_rate,
                            "episodes_per_second": episodes_per_second,
                            "active_workers": active,
                            "completion_percentage": completion_percentage,
                        }
                    )
                    self.logger.info(
                        f"Progress: {completed}/{total_work_items} work items completed "
                        f"({completion_percentage:.1f}%), {success}/{total} successful episodes "
                        f"({success_rate * 100:.1f}%), {active} workers active"
                    )
                    last_log_time = current_time
                except Exception as e:
                    self.logger.warning(f"WandB periodic logging failed: {e}")

            # Sleep briefly before checking again

            time.sleep(5)

        # Wait for all processes to complete
        for p in processes:
            p.join()
            p.close()

    else:
        # Single-worker mode runs in the main process
        house_processing_worker(
            worker_id=0,
            exp_config=self.config,
            work_items=self.work_items,
            shutdown_event=self.shutdown_event,
            counter_lock=self.counter_lock,
            house_counter=self.house_counter,
            success_count=self.success_count,
            total_count=self.total_count,
            completed_houses=self.completed_houses,
            skipped_houses=self.skipped_houses,
            max_allowed_sequential_task_sampler_failures=self.max_allowed_sequential_task_sampler_failures,
            max_allowed_sequential_rollout_failures=self.max_allowed_sequential_rollout_failures,
            max_allowed_sequential_irrecoverable_failures=self.max_allowed_sequential_irrecoverable_failures,
            preloaded_policy=preloaded_policy,
            filter_for_successful_trajectories=self.config.filter_for_successful_trajectories,
            runner_class=type(self),
        )

    # Extract final values from shared multiprocessing state
    success_count_val = self.success_count.value
    total_count_val = self.total_count.value
    completed_houses_val = self.completed_houses.value
    skipped_houses_val = self.skipped_houses.value

    success_rate = success_count_val / total_count_val if total_count_val > 0 else 0.0
    self.logger.info(
        f"Completed {completed_houses_val} work items, skipped {skipped_houses_val} work items"
    )
    self.logger.info(f"Success count: {success_count_val}, Total count: {total_count_val}")
    self.logger.info(f"Success rate: {success_rate * 100:.2f}%")

    # Log final metrics to WandB
    if self.wandb_enabled:
        try:
            final_elapsed_time = time.time() - start_time
            wandb.log(
                {
                    "final_success_count": success_count_val,
                    "final_total_count": total_count_val,
                    "final_success_rate": success_rate,
                    "final_completed_houses": completed_houses_val,
                    "final_skipped_houses": skipped_houses_val,
                    "final_elapsed_time_seconds": final_elapsed_time,
                    "final_elapsed_time_hours": final_elapsed_time / 3600,
                }
            )
            wandb.finish()
            self.logger.info("WandB logging finished")
        except Exception as e:
            self.logger.warning(f"WandB final logging failed: {e}")

    return success_count_val, total_count_val

run_single_rollout staticmethod

run_single_rollout(episode_seed: int, task: BaseMujocoTask, policy: Any, profiler: Profiler | None = None, viewer=None, shutdown_event=None, datagen_profiler: DatagenProfiler | None = None, end_on_success: bool = False) -> bool

Execute a single rollout with the given task and policy.

Parameters:

Name	Type	Description	Default
episode_seed	int	Seed for this episode	required
task	BaseMujocoTask	The task to run	required
policy	Any	Policy to use for action selection	required
profiler	Profiler | None	Legacy Profiler instance (optional)	None
viewer		MuJoCo viewer for visualization (optional)	None
shutdown_event		Event to signal shutdown (optional)	None
datagen_profiler	DatagenProfiler | None	DatagenProfiler for per-worker timing (optional)	None

Returns:

Name	Type	Description
bool	bool	Whether the episode was successful

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def run_single_rollout(
    episode_seed: int,
    task: BaseMujocoTask,
    policy: Any,
    profiler: Profiler | None = None,
    viewer=None,
    shutdown_event=None,
    datagen_profiler: DatagenProfiler | None = None,
    end_on_success: bool = False,
) -> bool:
    """Execute a single rollout with the given task and policy.

    Args:
        episode_seed: Seed for this episode
        task: The task to run
        policy: Policy to use for action selection
        profiler: Legacy Profiler instance (optional)
        viewer: MuJoCo viewer for visualization (optional)
        shutdown_event: Event to signal shutdown (optional)
        datagen_profiler: DatagenProfiler for per-worker timing (optional)

    Returns:
        bool: Whether the episode was successful
    """
    if profiler is not None:
        profiler.start("rollout")
    if datagen_profiler is not None:
        datagen_profiler.start("rollout_total")
        datagen_profiler.start("rollout_reset")

    observation, _info = task.reset()

    if datagen_profiler is not None:
        datagen_profiler.end("rollout_reset")

    if viewer is not None:
        viewer.sync()

    try:
        task.env.current_model.opt.enableflags |= int(mujoco.mjtEnableBit.mjENBL_SLEEP)
    except AttributeError:
        print("Not setting mujoco sleep. Needs version >=mujoco-3.8")

    step_count = 0
    while not task.is_done():
        # Check for shutdown signal
        if shutdown_event is not None and shutdown_event.is_set():
            if datagen_profiler is not None:
                datagen_profiler.end("rollout_total")
            return False

        # Step with policy
        if profiler is not None:
            profiler.start("policy_get_action")
        if datagen_profiler is not None:
            datagen_profiler.start("policy_get_action")
        action_cmd = policy.get_action(observation)
        if profiler is not None:
            profiler.end("policy_get_action")
        if datagen_profiler is not None:
            datagen_profiler.end("policy_get_action")

        # Step the task
        if profiler is not None:
            profiler.start("task_step")
        if datagen_profiler is not None:
            datagen_profiler.start("task_step")
        if action_cmd is None:
            print("Policy returned None action, ending episode")
            break
        observation, reward, terminal, truncated, infos = task.step(action_cmd)
        if profiler is not None:
            profiler.end("task_step")
        if datagen_profiler is not None:
            datagen_profiler.end("task_step")

        step_count += 1
        # Add termination if succ
        if end_on_success and "success" in infos[0] and infos[0]["success"]:
            success = True
            break

        if viewer is not None:
            viewer.sync()

    try:
        task.env.current_model.opt.enableflags &= ~int(mujoco.mjtEnableBit.mjENBL_SLEEP)
    except AttributeError:
        print("Not setting mujoco sleep. Needs version >=mujoco-3.8")

    # Save profiler summary
    if profiler is not None:
        profiler.end("rollout")
    if datagen_profiler is not None:
        datagen_profiler.end("rollout_total")
        # Record step count for reference
        datagen_profiler.record(
            "step_count_indicator", step_count / 1000.0
        )  # Scale down to avoid confusion

    # Check success if method exists
    success = task.judge_success() if hasattr(task, "judge_success") else False

    return success

sample_task_from_spec staticmethod

sample_task_from_spec(task_sampler, house_id: int, episode_spec, episode_idx: int) -> BaseMujocoTask | None

Sample task from specification.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def sample_task_from_spec(
    task_sampler, house_id: int, episode_spec, episode_idx: int
) -> BaseMujocoTask | None:
    """Sample task from specification."""
    return task_sampler.sample_task(house_index=house_id)

should_close_episode_task_sampler staticmethod

should_close_episode_task_sampler() -> bool

Whether to close task sampler after each episode.

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def should_close_episode_task_sampler() -> bool:
    """Whether to close task sampler after each episode."""
    return False

should_stop_early staticmethod

should_stop_early(num_collected: int, samples_per_house: int, exp_config: MlSpacesExpConfig | None = None) -> bool

Whether to stop before processing all episodes (e.g., enough successes).

Source code in molmo_spaces/data_generation/pipeline.py

@staticmethod
def should_stop_early(
    num_collected: int, samples_per_house: int, exp_config: "MlSpacesExpConfig | None" = None
) -> bool:
    """Whether to stop before processing all episodes (e.g., enough successes)."""
    return num_collected >= samples_per_house

cleanup_context

cleanup_context()

Context manager to ensure proper cleanup of MuJoCo resources

Source code in molmo_spaces/data_generation/pipeline.py

@contextmanager
def cleanup_context():
    """Context manager to ensure proper cleanup of MuJoCo resources"""
    try:
        yield
    finally:
        # Force garbage collection multiple times (CPython sometimes needs this)
        gc.collect()
        gc.collect()

        # Clear CUDA cache if available
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

cleanup_episode_resources

cleanup_episode_resources(task, policy, task_sampler, preloaded_policy: BasePolicy | None, close_task_sampler: bool = False) -> None

Cleanup resources after an episode.

Parameters:

Name	Type	Description	Default
task		Task instance to cleanup	required
policy		Policy instance to cleanup	required
task_sampler		Task sampler instance (optional cleanup)	required
preloaded_policy	BasePolicy | None	If not None, policy won't be deleted	required
close_task_sampler	bool	Whether to close the task sampler	False

Source code in molmo_spaces/data_generation/pipeline.py

def cleanup_episode_resources(
    task,
    policy,
    task_sampler,
    preloaded_policy: "BasePolicy | None",
    close_task_sampler: bool = False,
) -> None:
    """
    Cleanup resources after an episode.

    Args:
        task: Task instance to cleanup
        policy: Policy instance to cleanup
        task_sampler: Task sampler instance (optional cleanup)
        preloaded_policy: If not None, policy won't be deleted
        close_task_sampler: Whether to close the task sampler
    """
    if task is not None:
        if hasattr(task, "close"):
            task.close()

    if policy is not None and preloaded_policy is None:
        if hasattr(policy, "close"):
            policy.close()

    if close_task_sampler and task_sampler is not None:
        task_sampler.close()

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

get_detailed_memory_info

get_detailed_memory_info()

Get detailed memory usage information

Source code in molmo_spaces/data_generation/pipeline.py

def get_detailed_memory_info():
    """Get detailed memory usage information"""
    process = psutil.Process()
    mem = process.memory_info()
    return {
        "rss": mem.rss / 1024 / 1024,  # RSS in MB
        "vms": mem.vms / 1024 / 1024,  # Virtual Memory in MB
        "percent": process.memory_percent(),
    }

get_process_memory

get_process_memory()

Get current memory usage of the process in MB

Source code in molmo_spaces/data_generation/pipeline.py

def get_process_memory():
    """Get current memory usage of the process in MB"""
    process = psutil.Process()
    return process.memory_info().rss / 1024 / 1024

house_processing_worker

house_processing_worker(worker_id: int, exp_config: MlSpacesExpConfig, work_items: list[tuple[int, int, int, int]], shutdown_event, counter_lock, house_counter, success_count, total_count, completed_houses, skipped_houses, max_allowed_sequential_task_sampler_failures: int = 10, max_allowed_sequential_rollout_failures: int = 10, max_allowed_sequential_irrecoverable_failures: int = 5, preloaded_policy: BasePolicy | None = None, filter_for_successful_trajectories: bool = False, runner_class=None)

Standalone worker function that processes work items sequentially from a shared counter.

Each work item is a (house_id, batch_samples, batch_num, total_batches) tuple. This function can be run in either a thread or a process. It continually fetches the next work item from a shared counter and processes it.

Parameters:

Name	Type	Description	Default
worker_id	int	Unique ID for this worker	required
exp_config	MlSpacesExpConfig	Experiment configuration	required
work_items	list[tuple[int, int, int, int]]	List of (house_id, batch_samples, batch_num, total_batches) tuples	required
shutdown_event		Event to signal shutdown	required
counter_lock		Lock for thread-safe counter access	required
house_counter		Shared counter for next work item to process	required
success_count		Shared counter for successful episodes	required
total_count		Shared counter for total episodes	required
completed_houses		Shared counter for completed work items	required
skipped_houses		Shared counter for skipped work items	required
max_allowed_sequential_task_sampler_failures	int	Max consecutive task sampling failures	10
max_allowed_sequential_rollout_failures	int	Max consecutive rollout failures	10
max_allowed_sequential_irrecoverable_failures	int	Max consecutive irrecoverable failures	5
preloaded_policy	BasePolicy | None	Optional pre-initialized policy instance	None
filter_for_successful_trajectories	bool	Whether to filter for successful trajectories only	False
runner_class		Runner class with run_single_rollout and process_single_house static methods	None

Source code in molmo_spaces/data_generation/pipeline.py

def house_processing_worker(
    worker_id: int,
    exp_config: MlSpacesExpConfig,
    work_items: list[tuple[int, int, int, int]],
    shutdown_event,
    counter_lock,
    house_counter,
    success_count,
    total_count,
    completed_houses,
    skipped_houses,
    max_allowed_sequential_task_sampler_failures: int = 10,
    max_allowed_sequential_rollout_failures: int = 10,
    max_allowed_sequential_irrecoverable_failures: int = 5,
    preloaded_policy: BasePolicy | None = None,
    filter_for_successful_trajectories: bool = False,
    runner_class=None,
):
    """
    Standalone worker function that processes work items sequentially from a shared counter.

    Each work item is a (house_id, batch_samples, batch_num, total_batches) tuple.
    This function can be run in either a thread or a process. It continually fetches
    the next work item from a shared counter and processes it.

    Args:
        worker_id: Unique ID for this worker
        exp_config: Experiment configuration
        work_items: List of (house_id, batch_samples, batch_num, total_batches) tuples
        shutdown_event: Event to signal shutdown
        counter_lock: Lock for thread-safe counter access
        house_counter: Shared counter for next work item to process
        success_count: Shared counter for successful episodes
        total_count: Shared counter for total episodes
        completed_houses: Shared counter for completed work items
        skipped_houses: Shared counter for skipped work items
        max_allowed_sequential_task_sampler_failures: Max consecutive task sampling failures
        max_allowed_sequential_rollout_failures: Max consecutive rollout failures
        max_allowed_sequential_irrecoverable_failures: Max consecutive irrecoverable failures
        preloaded_policy: Optional pre-initialized policy instance
        filter_for_successful_trajectories: Whether to filter for successful trajectories only
        runner_class: Runner class with run_single_rollout and process_single_house static methods
    """
    # Create worker-specific logger
    worker_logger = get_worker_logger(worker_id)

    # Create per-worker profiler for timing analysis
    if hasattr(exp_config, "datagen_profiler") and exp_config.datagen_profiler:
        datagen_profiler = DatagenProfiler(logger=worker_logger, enabled=True)
    else:
        datagen_profiler = None

    # Track sequential irrecoverable failures at worker level
    num_sequential_irrecoverable_failures = 0

    # Normal datagen: create task sampler once for this worker (persists across all houses)
    # This allows the worker to track object diversity and other state across houses
    task_sampler = exp_config.task_sampler_config.task_sampler_class(exp_config)
    # Set profiler on task sampler for sub-timing within sample_task
    task_sampler.set_datagen_profiler(datagen_profiler)

    # Use context manager for worker-specific stdout redirection
    with worker_stdout_context(worker_logger, worker_id):
        try:
            while True:
                # Check for shutdown signal
                if shutdown_event.is_set():
                    worker_logger.info(
                        f"Worker {worker_id} received shutdown signal, cleaning up..."
                    )
                    break

                # Get next work item to process atomically
                with counter_lock:
                    if house_counter.value >= len(work_items):
                        break  # No more work items to process
                    item_idx = house_counter.value
                    house_counter.value += 1
                current_house_id, batch_samples, batch_num, total_batches = work_items[item_idx]

                worker_logger.info(
                    f"Worker {worker_id} starting house {current_house_id} "
                    f"batch {batch_num}/{total_batches} ({batch_samples} episodes) "
                    f"(item {item_idx}/{len(work_items)})"
                )

                # Process this work item
                house_success_count, house_total_count, irrecoverable = (
                    runner_class.process_single_house(
                        worker_id,
                        worker_logger,
                        current_house_id,
                        exp_config,
                        batch_samples,
                        shutdown_event,
                        task_sampler,
                        preloaded_policy,
                        max_allowed_sequential_task_sampler_failures,
                        max_allowed_sequential_rollout_failures,
                        filter_for_successful_trajectories=filter_for_successful_trajectories,
                        runner_class=runner_class,
                        batch_num=batch_num,
                        total_batches=total_batches,
                        datagen_profiler=datagen_profiler,
                    )
                )

                # Update global counters
                with counter_lock:
                    success_count.value += house_success_count
                    total_count.value += house_total_count
                    if house_total_count > 0:
                        completed_houses.value += 1
                    else:
                        skipped_houses.value += 1

                # Track sequential irrecoverable failures
                if irrecoverable:
                    num_sequential_irrecoverable_failures += 1
                    if (
                        num_sequential_irrecoverable_failures
                        >= max_allowed_sequential_irrecoverable_failures
                    ):
                        worker_logger.error(
                            f"Worker {worker_id} encountered {num_sequential_irrecoverable_failures} "
                            "sequential irrecoverable failures. This suggests something is seriously wrong. Exiting worker."
                        )
                        break
                else:
                    # Reset counter on success
                    num_sequential_irrecoverable_failures = 0

            worker_logger.info(f"Worker {worker_id} completed processing assigned work items")
        finally:
            # Log final profiling summary for this worker
            if datagen_profiler is not None:
                datagen_profiler.log_worker_summary()

            # Clean up task sampler at end of worker lifecycle
            if task_sampler is not None:
                task_sampler.close()

log_detailed_memory

log_detailed_memory(logger, prefix='') -> None

Log detailed memory usage information

Source code in molmo_spaces/data_generation/pipeline.py

def log_detailed_memory(logger, prefix="") -> None:
    """Log detailed memory usage information"""
    mem_info = get_detailed_memory_info()
    logger.info(f"{prefix}Memory Usage:")
    logger.info(f"  RSS: {mem_info['rss']:.1f} MB")
    logger.info(f"  Virtual: {mem_info['vms']:.1f} MB")
    logger.info(f"  Percent: {mem_info['percent']:.1f}%")

    # Log system memory
    system_mem = psutil.virtual_memory()
    logger.info("System Memory:")
    logger.info(f"  Total: {system_mem.total / 1024 / 1024:.1f} MB")
    logger.info(f"  Available: {system_mem.available / 1024 / 1024:.1f} MB")
    logger.info(f"  Used: {system_mem.used / 1024 / 1024:.1f} MB")
    logger.info(f"  Percent: {system_mem.percent:.1f}%")

log_memory_usage

log_memory_usage(logger, prefix='') -> None

Log current memory usage

Source code in molmo_spaces/data_generation/pipeline.py

def log_memory_usage(logger, prefix="") -> None:
    """Log current memory usage"""
    mem_usage = get_process_memory()
    logger.info(f"{prefix}Memory usage: {mem_usage:.2f} MB")

save_house_trajectories

save_house_trajectories(worker_logger, house_raw_histories: list, house_output_dir: Path, exp_config: MlSpacesExpConfig, batch_suffix: str, datagen_profiler: DatagenProfiler | None = None, batch_num: int | None = None, total_batches: int | None = None) -> None

Batch and save trajectory data for a house.

Parameters:

Name	Type	Description	Default
worker_logger		Logger instance	required
house_raw_histories	list	List of episode info dicts with 'history' and 'sensor_suite'	required
house_output_dir	Path	Output directory path	required
exp_config	MlSpacesExpConfig	Experiment configuration	required
batch_suffix	str	Suffix for batch file naming	required
datagen_profiler	DatagenProfiler | None	Profiler for timing	None
batch_num	int | None	Batch number for logging	None
total_batches	int | None	Total batches for logging	None

Source code in molmo_spaces/data_generation/pipeline.py

def save_house_trajectories(
    worker_logger,
    house_raw_histories: list,
    house_output_dir: Path,
    exp_config: "MlSpacesExpConfig",
    batch_suffix: str,
    datagen_profiler: "DatagenProfiler | None" = None,
    batch_num: int | None = None,
    total_batches: int | None = None,
) -> None:
    """
    Batch and save trajectory data for a house.

    Args:
        worker_logger: Logger instance
        house_raw_histories: List of episode info dicts with 'history' and 'sensor_suite'
        house_output_dir: Output directory path
        exp_config: Experiment configuration
        batch_suffix: Suffix for batch file naming
        datagen_profiler: Profiler for timing
        batch_num: Batch number for logging
        total_batches: Total batches for logging
    """
    if not house_raw_histories:
        worker_logger.warning(f"No trajectory data to save for {house_output_dir.name}")
        return

    batch_info = f" batch {batch_num}/{total_batches}" if batch_num is not None else ""
    worker_logger.info(
        f"Batching and saving trajectory data for {house_output_dir.name}{batch_info}: "
        f"{len(house_raw_histories)} episodes"
    )

    os.makedirs(house_output_dir, exist_ok=True)

    try:
        t_start = time.perf_counter()
        if datagen_profiler is not None:
            datagen_profiler.start("save_batch_prep")

        house_trajectory_data = []
        for idx, episode_info in enumerate(house_raw_histories):
            prepared_episode = prepare_episode_for_saving(
                episode_info["history"],
                episode_info["sensor_suite"],
                fps=exp_config.fps,
                save_dir=house_output_dir,
                episode_idx=idx,
                save_file_suffix=batch_suffix,
            )
            if prepared_episode is not None:
                house_trajectory_data.append(prepared_episode)
            del episode_info["history"]

        t_batch = time.perf_counter() - t_start
        if datagen_profiler is not None:
            datagen_profiler.end("save_batch_prep")
        worker_logger.debug(f"Batched {len(house_trajectory_data)} episodes in {t_batch:.2f}s")

        t_save_start = time.perf_counter()
        if datagen_profiler is not None:
            datagen_profiler.start("save_trajectories")
        save_trajectories(
            house_trajectory_data,
            save_dir=house_output_dir,
            fps=exp_config.fps,
            save_file_suffix=batch_suffix,
            save_mp4s=True,
            logger=worker_logger,
        )
        if datagen_profiler is not None:
            datagen_profiler.end("save_trajectories")
        t_save = time.perf_counter() - t_save_start

        total_time = time.perf_counter() - t_start
        worker_logger.info(
            f"Successfully saved trajectory data for {house_output_dir.name} in {total_time:.2f}s "
            f"(batch: {t_batch:.2f}s, save: {t_save:.2f}s)"
        )

        del house_trajectory_data
        gc.collect()

    except Exception as e:
        worker_logger.error(f"Failed to save trajectory data for {house_output_dir.name}: {e}")
        traceback.print_exc()

setup_house_dirs

setup_house_dirs(exp_config: MlSpacesExpConfig, house_id: int, batch_num: int | None = None, total_batches: int | None = None) -> tuple[Path, Path, str, bool]

Setup output directories and check for existing output.

Parameters:

Name	Type	Description	Default
exp_config	MlSpacesExpConfig	Experiment configuration	required
house_id	int	House index	required
batch_num	int | None	Batch number (1-indexed)	None
total_batches	int | None	Total number of batches	None

Returns:

Name	Type	Description
tuple	tuple[Path, Path, str, bool]	(house_output_dir, house_debug_dir, batch_suffix, should_skip)

Source code in molmo_spaces/data_generation/pipeline.py

def setup_house_dirs(
    exp_config: "MlSpacesExpConfig",
    house_id: int,
    batch_num: int | None = None,
    total_batches: int | None = None,
) -> tuple[Path, Path, str, bool]:
    """
    Setup output directories and check for existing output.

    Args:
        exp_config: Experiment configuration
        house_id: House index
        batch_num: Batch number (1-indexed)
        total_batches: Total number of batches

    Returns:
        tuple: (house_output_dir, house_debug_dir, batch_suffix, should_skip)
    """
    house_output_dir = exp_config.output_dir / f"house_{house_id}"
    debug_base_dir = exp_config.output_dir.parent / "debug" / exp_config.output_dir.name
    house_debug_dir = debug_base_dir / f"house_{house_id}"

    if batch_num is None or total_batches is None:
        batch_num = 1
        total_batches = 1
    batch_suffix = f"_batch_{batch_num}_of_{total_batches}"

    batch_file = house_output_dir / f"trajectories{batch_suffix}.h5"
    should_skip = batch_file.exists()

    return house_output_dir, house_debug_dir, batch_suffix, should_skip

setup_policy

setup_policy(exp_config: MlSpacesExpConfig, task: BaseMujocoTask, preloaded_policy: BasePolicy | None, datagen_profiler: DatagenProfiler | None) -> BasePolicy

Create or return policy for episode.

Parameters:

Name	Type	Description	Default
exp_config	MlSpacesExpConfig	Experiment configuration	required
task	BaseMujocoTask	The task instance	required
preloaded_policy	BasePolicy | None	Pre-loaded policy or None	required
datagen_profiler	DatagenProfiler | None	Profiler for timing	required

Returns:

Type	Description
BasePolicy	Policy instance

Source code in molmo_spaces/data_generation/pipeline.py

def setup_policy(
    exp_config: "MlSpacesExpConfig",
    task: "BaseMujocoTask",
    preloaded_policy: "BasePolicy | None",
    datagen_profiler: "DatagenProfiler | None",
) -> "BasePolicy":
    """
    Create or return policy for episode.

    Args:
        exp_config: Experiment configuration
        task: The task instance
        preloaded_policy: Pre-loaded policy or None
        datagen_profiler: Profiler for timing

    Returns:
        Policy instance
    """
    if datagen_profiler is not None:
        datagen_profiler.start("policy_setup")

    if preloaded_policy is not None:
        policy = preloaded_policy
    else:
        policy = exp_config.policy_config.policy_factory(exp_config, task)

    task.register_policy(policy)

    if datagen_profiler is not None:
        datagen_profiler.end("policy_setup")

    return policy

setup_viewer

setup_viewer(exp_config: MlSpacesExpConfig, task: BaseMujocoTask, policy: BasePolicy, current_viewer)

Setup passive viewer if configured.

Parameters:

Name	Type	Description	Default
exp_config	MlSpacesExpConfig	Experiment configuration	required
task	BaseMujocoTask	The task instance	required
policy	BasePolicy	The policy instance	required
current_viewer		Existing viewer or None	required

Returns:

Type	Description
	Viewer instance or None

Source code in molmo_spaces/data_generation/pipeline.py

def setup_viewer(
    exp_config: "MlSpacesExpConfig",
    task: "BaseMujocoTask",
    policy: "BasePolicy",
    current_viewer,
):
    """
    Setup passive viewer if configured.

    Args:
        exp_config: Experiment configuration
        task: The task instance
        policy: The policy instance
        current_viewer: Existing viewer or None

    Returns:
        Viewer instance or None
    """
    viewer = current_viewer
    if exp_config.use_passive_viewer:
        if viewer is not None:
            viewer.close()
            viewer = None
        import mujoco.viewer

        viewer = mujoco.viewer.launch_passive(
            task.env.mj_datas[task.env.current_batch_index].model,
            task.env.mj_datas[task.env.current_batch_index],
            key_callback=getattr(policy, "get_key_callback", lambda: None)(),
        )
        if exp_config.viewer_cam_dict is not None and "camera" in exp_config.viewer_cam_dict:
            viewer.cam.fixedcamid = (
                task.env.mj_datas[0].camera(exp_config.viewer_cam_dict["camera"]).id
            )
            viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
        viewer.opt.sitegroup[0] = False
    task.viewer = viewer
    return viewer