__init__.py

"""Robot control stack python bindings."""

import os
import site
from dataclasses import dataclass, field

from gymnasium import register
from lerobot.configs.types import PolicyFeature
from lerobot.envs.configs import EnvConfig
from rcs._core import __version__, common
from rcs.envs.creators import (
    FR3LabDigitGripperPickUpSimEnvCreator,
    FR3SimplePickUpSimEnvCreator,
)
from rcs.lerobot import Pick

from rcs import camera, envs, hand, sim


@dataclass(kw_only=True)
class Scene:
    """Scene configuration."""

    mjb: str
    """Path to the Mujoco binary scene file."""
    mjcf_scene: str
    """Path to the Mujoco scene XML file."""
    mjcf_robot: str
    """Path to the Mujoco robot XML file for IK."""
    urdf: str | None = None
    """Path to the URDF robot file for IK, if available."""
    robot_type: common.RobotType
    """Type of the robot in the scene."""


def get_scene_urdf(scene_name: str) -> str | None:
    urdf_path = os.path.join(site.getsitepackages()[0], "rcs", "scenes", scene_name, "robot.urdf")
    return urdf_path if os.path.exists(urdf_path) else None


# available mujoco scenes
scenes: dict[str, Scene] = {
    "fr3_empty_world": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_empty_world", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_empty_world", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_empty_world", "robot.xml"),
        urdf=get_scene_urdf("fr3_empty_world"),
        robot_type=common.RobotType.FR3,
    ),
    "fr3_simple_pick_up": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_simple_pick_up", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_simple_pick_up", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_simple_pick_up", "robot.xml"),
        urdf=get_scene_urdf("fr3_simple_pick_up"),
        robot_type=common.RobotType.FR3,
    ),
    "fr3_digit_simple_pick_up": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_digit_simple_pick_up", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_digit_simple_pick_up", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_digit_simple_pick_up", "fr3_0.xml"),
        urdf=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "fr3_digit_simple_pick_up", "robot.urdf"),
        robot_type=common.RobotType.FR3,
    ),
    "xarm7_empty_world": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "xarm7_empty_world", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "xarm7_empty_world", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "xarm7_empty_world", "robot.xml"),
        urdf=get_scene_urdf("xarm7_empty_world"),
        robot_type=common.RobotType.XArm7,
    ),
    "panda_empty_world": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "panda_empty_world", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "panda_empty_world", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "panda_empty_world", "robot.xml"),
        robot_type=common.RobotType.Panda,
    ),
    "so101_empty_world": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "so101_empty_world", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "so101_empty_world", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "so101_empty_world", "robot.xml"),
        urdf=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "so101_empty_world", "robot.urdf"),
        robot_type=common.RobotType.SO101,
    ),
    "ur5e_empty_world": Scene(
        mjb=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "ur5e_empty_world", "scene.mjb"),
        mjcf_scene=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "ur5e_empty_world", "scene.xml"),
        mjcf_robot=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "ur5e_empty_world", "robot.xml"),
        urdf=os.path.join(site.getsitepackages()[0], "rcs", "scenes", "ur5e_empty_world", "robot.urdf"),
        robot_type=common.RobotType.UR5e,
    ),
}

# make submodules available
__all__ = ["__doc__", "__version__", "common", "sim", "camera", "scenes", "envs", "hand"]

# register gymnasium environments
register(
    id="rcs/FR3SimplePickUpSim-v0",
    entry_point=FR3SimplePickUpSimEnvCreator(),
)
register(
    id="rcs/FR3LabDigitGripperPickUpSim-v0",
    entry_point=FR3LabDigitGripperPickUpSimEnvCreator(),
)
register(
    id="rcs/pick",
    entry_point=Pick(),
)

# TODO: gym.make("rcs/FR3SimEnv-v0") results in a pickling error:
# TypeError: cannot pickle 'rcs._core.sim.SimRobotConfig' object
# cf. https://pybind11.readthedocs.io/en/stable/advanced/classes.html#deepcopy-support
# register(
#    id="rcs/FR3SimEnv-v0",
#    entry_point=SimEnvCreator(),
# )


# lerobot env registration
@EnvConfig.register_subclass("rcs/pick")
@dataclass
class RCSEnvConfig(EnvConfig):
    task: str | None = "pick"
    fps: int = 30
    features: dict[str, PolicyFeature] = field(default_factory=dict)
    features_map: dict[str, str] = field(default_factory=dict)
    max_parallel_tasks: int = 1
    disable_env_checker: bool = True

    @property
    def package_name(self) -> str:
        """Package name to import if environment not found in gym registry"""
        return "rcs"

    @property
    def gym_id(self) -> str:
        """ID string used in gym.make() to instantiate the environment"""
        return f"{self.package_name}/{self.task}"

    @property
    def gym_kwargs(self) -> dict:
        return {}