molmo-spaces 0.2.9__py3-none-any.whl

molmo_spaces/__init__.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+MOLMO_SPACES_BASE_DIR = Path(__file__).parent

molmo_spaces/configs/__init__.py

@@ -0,0 +1,61 @@
+"""Configuration module for MolmoSpaces experiments.
+
+This module provides configuration classes organized by category:
+- abstract_config: Base Config class
+- abstract_exp_config: Base experiment configuration
+- camera_configs: Camera-related configurations
+- robot_configs: Robot-related configurations
+- task_configs: Task-related configurations
+- task_sampler_configs: Task sampler-related configurations
+- policy_configs: Policy-related configurations
+"""
+
+from molmo_spaces.configs.abstract_config import Config
+from molmo_spaces.configs.abstract_exp_config import MlSpacesExpConfig
+from molmo_spaces.configs.camera_configs import (
+    CameraConfig,
+    CameraSystemConfig,
+    FixedExocentricCameraConfig,
+    FrankaDroidCameraSystem,
+    FrankaRandomizedD405D455CameraSystem,
+    MjcfCameraConfig,
+    RandomizedExocentricCameraConfig,
+    RBY1GoProD455CameraSystem,
+    RBY1MjcfCameraSystem,
+    RobotMountedCameraConfig,
+)
+from molmo_spaces.configs.policy_configs import BasePolicyConfig
+from molmo_spaces.configs.robot_configs import BaseRobotConfig, FrankaRobotConfig
+from molmo_spaces.configs.task_configs import BaseMujocoTaskConfig, PickTaskConfig
+from molmo_spaces.configs.task_sampler_configs import (
+    BaseMujocoTaskSamplerConfig,
+    PickTaskSamplerConfig,
+)
+
+__all__ = [
+    "Config",
+    "MlSpacesExpConfig",
+    # Camera configs - new unified system
+    "CameraSystemConfig",
+    "CameraConfig",
+    "MjcfCameraConfig",
+    "RobotMountedCameraConfig",
+    "FixedExocentricCameraConfig",
+    "RandomizedExocentricCameraConfig",
+    "RBY1MjcfCameraSystem",
+    "RBY1GoProD455CameraSystem",
+    "FrankaRandomizedD405D455CameraSystem",
+    "FrankaDroidCameraSystem",
+    # Robot configs
+    "BaseRobotConfig",
+    "FrankaRobotConfig",
+    # Task configs
+    "BaseMujocoTaskConfig",
+    "PickTaskConfig",
+    # Task sampler configs
+    "BaseMujocoTaskSamplerConfig",
+    "PickTaskSamplerConfig",
+    # Policy configs
+    "BasePolicyConfig",
+    "ObjectManipulationPlannerPolicyConfig",
+]

molmo_spaces/configs/abstract_config.py

@@ -0,0 +1,43 @@
+"""Simple configuration management using Pydantic.
+This module provides a base configuration class that can be extended to create specific configurations.
+It uses Pydantic for data validation and can return dicts or jsons or save or load jsons from files.
+"""
+
+from pydantic import BaseModel, ConfigDict
+
+
+class Config(BaseModel):
+    """
+    Base configuration class that can be extended for specific configurations.
+    Provides methods to convert to dict, json, and to save/load from files.
+    """
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,  # allow arbitrary types in the config
+        # frozen=True,  # (Optional) Make the config immutable
+    )
+
+    def to_dict(self) -> dict:
+        """Convert the configuration to a dictionary."""
+        return self.model_dump()
+
+    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")
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Config":
+        """Create a configuration instance from a dictionary."""
+        return cls.model_validate(data)
+
+    @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)
+
+    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())

molmo_spaces/configs/abstract_exp_config.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import base64
+import datetime
+import logging
+import pickle
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, ClassVar
+
+from molmo_spaces.configs.abstract_config import Config
+from molmo_spaces.configs.camera_configs import (
+    AllCameraSystems,
+    CameraSystemConfig,
+    FixedExocentricCameraConfig,
+    MjcfCameraConfig,
+    RandomizedExocentricCameraConfig,
+    RobotMountedCameraConfig,
+)
+from molmo_spaces.configs.policy_configs import BasePolicyConfig
+from molmo_spaces.configs.robot_configs import BaseRobotConfig
+from molmo_spaces.configs.task_configs import AllTaskConfigs
+from molmo_spaces.configs.task_sampler_configs import BaseMujocoTaskSamplerConfig
+from molmo_spaces.tasks.task import BaseMujocoTask
+from molmo_spaces.utils.pose import pose_mat_to_7d
+from molmo_spaces.utils.profiler_utils import Profiler
+
+log = logging.getLogger(__name__)
+
+
+class MlSpacesExpConfig(Config, ABC):
+    """
+    Base configuration class for experiments.
+    This should be extended to create specific experiment configurations.
+    """
+
+    config_version: str = "0.1"
+    num_envs: (
+        int  # Number of batched environments per worker (for vectorized physics in CPUMujocoEnv)
+    )
+    num_workers: int = (
+        1  # Number of worker processes for parallel data generation (episode-level parallelism)
+    )
+
+    task_type: str  # Task type: e.g. pick, pick_and_place, etc.
+    use_passive_viewer: bool  # Launch passive viewer for rendering
+    viewer_cam_dict: dict  # Dictionary containing viewer camera parameters
+    policy_dt_ms: float  # Default policy time step
+    ctrl_dt_ms: float  # Default control time step
+    sim_dt_ms: float  # Default simulation time step
+    seed: int | None = None  # Random seed for task sampling (if None, generates random seed)
+    task_horizon: int | None = None  # Maximum number of steps per episode (if None, no time limit)
+    end_on_success: bool = (
+        False  # Whether to end episode immediately upon success (overrides task_horizon if True)
+    )
+    collision_free_pose_limit: int = 3
+
+    # Scene configuration
+    scene_dataset: str  # Scenes to use, e.g. ithor, procthor-10k, procthor-objaverse. If "user", use the scene_xml_paths in task_sampler_config.
+    data_split: str = "train"  # Data split to use, e.g. train, val, test
+
+    @property
+    def fps(self) -> float:
+        return 1000.0 / self.policy_dt_ms
+
+    # Configuration fields using imported base classes
+    camera_config: CameraSystemConfig | None = None  # Configuration for cameras and sensors
+    robot_config: BaseRobotConfig  # Configuration for the robot
+    task_sampler_config: BaseMujocoTaskSamplerConfig  # Configuration for the task sampler
+    task_config: AllTaskConfigs  # Configuration for tasks
+    task_config_preset_exp: AllTaskConfigs | None = None  # Cached config for whole experiment
+    task_config_preset_scn: AllTaskConfigs | None = None  # Cached config for scene
+    policy_config: BasePolicyConfig  # Configuration for policies
+
+    benchmark_path: Path | None = None  # Contains a json with a list of fully-specified episodes.
+
+    # Evaluation runtime parameters (optional, set during evaluation initialization)
+    # Using Any to avoid circular import - actual type is EvalRuntimeParams from eval_main
+    eval_runtime_params: Any = None
+
+    # Output and profiling
+    output_dir: Path  # Output directory for experiment results
+    profile: bool = False  # Whether to enable profiling
+    profiler: Profiler | None = None  # Profiler instance (auto-created if profile=True)
+    datagen_profiler: bool = True  # run the datagen profiler
+
+    # Logging
+    log_level: str = "info"  # Global logging level: "debug", "info", "warning", "error", "none"
+    use_wandb: bool = False
+    wandb_project: str | None = None
+    wandb_name: str | None = None  # Weights & Biases run name
+
+    # Backward compatibility aliases for nested class references (ClassVar so Pydantic ignores them)
+    CameraConfig: ClassVar[type] = CameraSystemConfig
+    RobotConfig: ClassVar[type] = BaseRobotConfig
+    PolicyConfig: ClassVar[type] = BasePolicyConfig
+
+    filter_for_successful_trajectories: bool = True  # If True, only save successful trajectories to main output directory (failed episodes may be sampled 1% for debug directory). If False, save all trajectories to main output directory.
+
+    environment_light_intensity: float = 15000.0
+
+    no_cached_map: bool = False
+
+    def model_post_init(self, _context) -> None:
+        """This serves as the __init__() called after internal validation of config parameters"""
+        assert (self.policy_dt_ms / self.ctrl_dt_ms).is_integer(), (
+            "policy_dt_ms must be a multiple of ctrl_dt_ms"
+        )
+        assert (self.ctrl_dt_ms / self.sim_dt_ms).is_integer(), (
+            "ctrl_dt_ms must be a multiple of sim_dt"
+        )
+
+    @property
+    @abstractmethod
+    def tag(self) -> str:
+        """A string describing the experiment."""
+
+    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}")
+
+    @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
+
+    # The code below this is used for saving episode state so that it can be re-loaded w/o sampling
+    class SavedEpisode(Config):
+        """Config informationd describing a sinlge episode"""
+
+        camera_config: AllCameraSystems | None = None  # Configuration for cameras and sensors
+        robot_config: BaseRobotConfig | None = None  # Configuration for the robot
+        task_config: AllTaskConfigs | None = None  # Configuration for tasks
+        task_cls_str: str | None = None
+
+    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

molmo_spaces/configs/base_nav_to_obj_config.py

@@ -0,0 +1,112 @@
+"""
+Example configuration for RBY1 navigation to object data generation using the extracted task sampler.
+This shows how the scene randomization functionality from the reference script has been
+properly integrated into the modular task sampler architecture.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from molmo_spaces.configs.abstract_config import Config
+from molmo_spaces.configs.abstract_exp_config import MlSpacesExpConfig
+from molmo_spaces.configs.camera_configs import RBY1MjcfCameraSystem
+from molmo_spaces.configs.policy_configs import AStarNavToObjPolicyConfig, BasePolicyConfig
+from molmo_spaces.configs.robot_configs import BaseRobotConfig, RBY1Config
+from molmo_spaces.configs.task_configs import NavToObjTaskConfig
+from molmo_spaces.configs.task_sampler_configs import NavToObjTaskSamplerConfig
+from molmo_spaces.tasks.nav_task import NavToObjTask
+from molmo_spaces.tasks.nav_task_sampler import NavToObjTaskSampler
+from molmo_spaces.utils.profiler_utils import Profiler
+
+
+class NavToObjBaseConfig(MlSpacesExpConfig):
+    """Base configuration for navigation to object data generation tasks."""
+
+    # NOTE: will not work if used directly. Subclass examples in data_generation/configs.py
+
+    # --- Experiment-level config parameters ---
+    num_envs: int = 1  # Number of environments to run in each thread
+    use_passive_viewer: bool = False  # Launch passive viewer for rendering
+    viewer_camera: None = None
+    viewer_cam_dict: dict = {
+        "distance": 5.0,
+        "azimuth": 45.0,
+        "elevation": -30.0,
+        "lookat": np.array([0.0, 0.0, 0.5]),
+    }
+    policy_dt_ms: float = 200.0  # policy time step
+    ctrl_dt_ms: float = 2.0  # control time step
+    sim_dt_ms: float = 2.0  # simulation time step
+    task_horizon: int = 500  # Maximum steps per episode to prevent infinite runs
+    record_videos: bool = False  # Whether to record videos of episodes
+
+    # --- Data generation settings ---
+    num_threads: int = 1  # parallel data generation threads
+    profile: bool = True  # Whether to profile the data generation pipeline
+    profiler: Profiler | None = None
+    output_dir: str | None = None  # Directory to save generated data
+    use_wandb: bool = False  # Whether to use Weights & Biases logging
+    wandb_name: str | None = None  # Weights & Biases run name
+    wandb_project: str | None = None  # Weights & Biases project name
+
+    # --- Task type configuration ---
+    task_type: str = "nav_to_obj"  # Task type: nav_to_obj
+
+    # --- ProcTHOR dataset configuration ---
+    scene_dataset: str = "procthor-10k"  # Name of the scene dataset to load
+    data_split: str = "train"  # Data split to use
+
+    robot_config: BaseRobotConfig | None = None
+
+    # Camera configuration - using new unified camera system
+    camera_config: RBY1MjcfCameraSystem = RBY1MjcfCameraSystem()
+
+    # Task sampler configuration (imported from task_sampler_configs.py)
+    task_sampler_config: NavToObjTaskSamplerConfig = NavToObjTaskSamplerConfig(
+        task_sampler_class=NavToObjTaskSampler
+    )
+
+    # Task configuration (imported from task_configs.py)
+    task_config: NavToObjTaskConfig = NavToObjTaskConfig(task_cls=NavToObjTask)
+    task_config_preset: NavToObjTaskConfig | None = None
+
+    # Policy configuration (imported from policy_configs.py)
+    policy_config: BasePolicyConfig = AStarNavToObjPolicyConfig()
+
+    def _init_policy_config(self) -> BasePolicyConfig:
+        """Initialize policy config. Override in subclasses for dynamic initialization."""
+        return self.policy_config
+
+    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()
+
+    @property
+    def tag(self) -> str:
+        return "nav_to_obj_datagen"
+
+    class SavedEpisode(Config):
+        camera_config: RBY1MjcfCameraSystem | None = None  # Configuration for cameras and sensors
+        robot_config: RBY1Config | None = None  # Configuration for the robot
+        task_config: NavToObjTaskConfig | None = None  # Configuration for tasks
+        task_cls_str: str | None = None

molmo_spaces/configs/base_open_task_configs.py

@@ -0,0 +1,56 @@
+from molmo_spaces.configs.base_pick_config import PickBaseConfig
+from molmo_spaces.configs.policy_configs import BasePolicyConfig, OpenClosePlannerPolicyConfig
+from molmo_spaces.configs.task_configs import OpeningTaskConfig
+from molmo_spaces.configs.task_sampler_configs import (
+    OpenTaskSamplerConfig,
+)
+from molmo_spaces.tasks.opening_task_samplers import OpenTaskSampler
+from molmo_spaces.tasks.opening_tasks import OpeningTask
+
+
+class OpeningBaseConfig(PickBaseConfig):
+    """Base configuration for opening task data generation."""
+
+    task_type: str = "open"
+
+    # Task sampler configuration (imported from task_sampler_configs.py)
+    task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(
+        task_sampler_class=OpenTaskSampler,
+        target_initial_state_open_percentage=0,  # 0 for open task
+    )
+
+    # Task configuration (imported from task_configs.py)
+    task_config: OpeningTaskConfig = OpeningTaskConfig(
+        task_cls=OpeningTask,
+        task_success_threshold=0.15,  # low for now, due to placement/IK constraints, should be ~0.66
+        joint_index=0,
+        any_inst_of_category=True,  # open any instance of category
+    )
+    task_config_preset: OpeningTaskConfig | None = None
+
+    # Policy configuration (imported from policy_configs.py)
+    policy_config: BasePolicyConfig = OpenClosePlannerPolicyConfig()
+
+
+class ClosingBaseConfig(PickBaseConfig):
+    """Base configuration for closing task data generation."""
+
+    task_type: str = "close"
+
+    # Task sampler configuration (imported from task_sampler_configs.py)
+    task_sampler_config: OpenTaskSamplerConfig = OpenTaskSamplerConfig(
+        task_sampler_class=OpenTaskSampler,
+        target_initial_state_open_percentage=0.5,  # 0.67 for close task
+    )
+
+    # Task configuration (imported from task_configs.py)
+    task_config: OpeningTaskConfig = OpeningTaskConfig(
+        task_cls=OpeningTask,
+        task_success_threshold=0.85,  # For closing Task, 0.33
+        joint_index=0,
+        any_inst_of_category=False,  # open any instance of category
+    )
+    task_config_preset: OpeningTaskConfig | None = None
+
+    # Policy configuration (imported from policy_configs.py)
+    policy_config: BasePolicyConfig = OpenClosePlannerPolicyConfig()

molmo_spaces/configs/base_packing_configs.py

@@ -0,0 +1,21 @@
+from molmo_spaces.configs.base_pick_config import PickBaseConfig
+from molmo_spaces.configs.policy_configs import PickAndPlacePlannerPolicyConfig
+from molmo_spaces.configs.task_configs import PackingTaskConfig
+from molmo_spaces.configs.task_sampler_configs import PackingTaskSamplerConfig
+from molmo_spaces.data_generation.config_registry import register_config
+from molmo_spaces.tasks.packing_task import PackingTask
+from molmo_spaces.tasks.packing_task_sampler import PackingTaskSampler
+from molmo_spaces.utils.constants.object_constants import PICK_AND_PLACE_OBJECTS
+
+
+@register_config("PackingDataGenConfig")
+class PackingDataGenConfig(PickBaseConfig):
+    task_type: str = "packing"
+    num_workers: int = 1
+    task_sampler_config: PackingTaskSamplerConfig = PackingTaskSamplerConfig(
+        task_sampler_class=PackingTaskSampler,
+        pickup_types=PICK_AND_PLACE_OBJECTS,
+        samples_per_house=20,
+    )
+    task_config: PackingTaskConfig = PackingTaskConfig(task_cls=PackingTask)
+    policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()

molmo_spaces/configs/base_pick_and_place_color_configs.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+
+from molmo_spaces.configs.base_pick_config import PickBaseConfig
+from molmo_spaces.configs.policy_configs import PickAndPlaceColorPlannerPolicyConfig
+from molmo_spaces.configs.task_configs import PickAndPlaceColorTaskConfig
+from molmo_spaces.configs.task_sampler_configs import PickAndPlaceColorTaskSamplerConfig
+from molmo_spaces.data_generation.config_registry import register_config
+from molmo_spaces.molmo_spaces_constants import ASSETS_DIR
+from molmo_spaces.tasks.pick_and_place_color_task import PickAndPlaceColorTask
+from molmo_spaces.tasks.pick_and_place_color_task_sampler import PickAndPlaceColorTaskSampler
+
+
+@register_config("PickAndPlaceColorDataGenConfig")
+class PickAndPlaceColorDataGenConfig(PickBaseConfig):
+    task_type: str = "pick_and_place_color"
+    num_workers: int = 1
+    output_dir: Path = ASSETS_DIR / "experiment_output" / "datagen" / "pick_and_place_color_base_v1"
+    wandb_project: str = "molmo-spaces-data-generation"
+    task_sampler_config: PickAndPlaceColorTaskSamplerConfig = PickAndPlaceColorTaskSamplerConfig(
+        task_sampler_class=PickAndPlaceColorTaskSampler,
+        samples_per_house=20,
+    )
+    task_config: PickAndPlaceColorTaskConfig = PickAndPlaceColorTaskConfig(
+        task_cls=PickAndPlaceColorTask
+    )
+    policy_config: PickAndPlaceColorPlannerPolicyConfig = PickAndPlaceColorPlannerPolicyConfig()
+    use_passive_viewer: bool = False
+
+    @property
+    def tag(self) -> str:
+        return "pick_and_place_color_base_datagen"

molmo_spaces/configs/base_pick_and_place_configs.py

@@ -0,0 +1,27 @@
+from molmo_spaces.configs.abstract_config import Config
+from molmo_spaces.configs.base_pick_config import PickBaseConfig
+from molmo_spaces.configs.camera_configs import AllCameraSystems
+from molmo_spaces.configs.policy_configs import PickAndPlacePlannerPolicyConfig
+from molmo_spaces.configs.robot_configs import FrankaRobotConfig
+from molmo_spaces.configs.task_configs import PickAndPlaceTaskConfig
+from molmo_spaces.configs.task_sampler_configs import PickAndPlaceTaskSamplerConfig
+from molmo_spaces.tasks.pick_and_place_task import PickAndPlaceTask
+from molmo_spaces.tasks.pick_and_place_task_sampler import PickAndPlaceTaskSampler
+
+
+class PickAndPlaceDataGenConfig(PickBaseConfig):
+    task_type: str = "pick_and_place"
+    num_workers: int = 1
+    task_sampler_config: PickAndPlaceTaskSamplerConfig = PickAndPlaceTaskSamplerConfig(
+        task_sampler_class=PickAndPlaceTaskSampler,
+        pickup_types=[],
+        samples_per_house=20,
+    )
+    task_config: PickAndPlaceTaskConfig = PickAndPlaceTaskConfig(task_cls=PickAndPlaceTask)
+    policy_config: PickAndPlacePlannerPolicyConfig = PickAndPlacePlannerPolicyConfig()
+
+    class SavedEpisode(Config):
+        camera_config: AllCameraSystems | None = None  # Configuration for cameras and sensors
+        robot_config: FrankaRobotConfig | None = None  # Configuration for the robot
+        task_config: PickAndPlaceTaskConfig | None = None  # Configuration for tasks
+        task_cls_str: str | None = None

molmo_spaces/configs/base_pick_and_place_next_to_configs.py

@@ -0,0 +1,26 @@
+from molmo_spaces.configs.base_pick_config import PickBaseConfig
+from molmo_spaces.configs.policy_configs import PickAndPlaceNextToPlannerPolicyConfig
+from molmo_spaces.configs.robot_configs import BaseRobotConfig, FrankaRobotConfig
+from molmo_spaces.configs.task_configs import PickAndPlaceNextToTaskConfig
+from molmo_spaces.configs.task_sampler_configs import PickAndPlaceNextToTaskSamplerConfig
+from molmo_spaces.data_generation.config_registry import register_config
+from molmo_spaces.tasks.pick_and_place_next_to_task import PickAndPlaceNextToTask
+from molmo_spaces.tasks.pick_and_place_next_to_task_sampler import PickAndPlaceNextToTaskSampler
+from molmo_spaces.utils.constants.object_constants import PICK_AND_PLACE_OBJECTS
+
+
+@register_config("PickAndPlaceNextToDataGenConfig")
+class PickAndPlaceNextToDataGenConfig(PickBaseConfig):
+    task_type: str = "pick_and_place_next_to"
+    num_workers: int = 1
+    task_sampler_config: PickAndPlaceNextToTaskSamplerConfig = PickAndPlaceNextToTaskSamplerConfig(
+        task_sampler_class=PickAndPlaceNextToTaskSampler,
+        pickup_types=PICK_AND_PLACE_OBJECTS,
+        samples_per_house=20,
+    )
+    task_config: PickAndPlaceNextToTaskConfig = PickAndPlaceNextToTaskConfig(
+        task_cls=PickAndPlaceNextToTask
+    )
+    policy_config: PickAndPlaceNextToPlannerPolicyConfig = PickAndPlaceNextToPlannerPolicyConfig()
+    robot_config: BaseRobotConfig = FrankaRobotConfig()
+    use_passive_viewer: bool = False