LADDER SimplerEnv-Long for LeRobot EnvHub (LADDER v1.0.0)

21fed8e verified 3 days ago

6.9 kB

# Copyright (c) 2026 INSAIT. MIT License (see LICENSE).
"""LADDER SimplerEnv-Long for LeRobot EnvHub: 54 long-horizon WidowX tasks on SAPIEN 3 / ManiSkill 3.
Each task is one suite of LeRobot's ``{suite: {task_id: env}}``; its 24 episodes are the 24 object
layouts LADDER defines (24 parallel envs on one GPU). See README.md.
"""

import os
import re
import sys

REPO_ID = "INSAIT-Institute/LADDER-SimplerEnv-Long"

_HERE = os.path.dirname(os.path.abspath(__file__))
if not os.path.exists(os.path.join(_HERE, "_ladder_envhub.py")):
    # `make_env("<repo>")` downloads only this file: fetch the rest of the same revision.
    from huggingface_hub import snapshot_download

    _revision = os.path.basename(_HERE)
    _HERE = snapshot_download(REPO_ID, revision=_revision if re.fullmatch(r"[0-9a-f]{40}", _revision) else None)
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

from functools import partial  # noqa: E402
from typing import Any  # noqa: E402

import _ladder_envhub as lh  # noqa: E402

EPISODES_PER_TASK = 24
# ManiSkill reads its asset directory once, when first imported; the tree is filled on the first reset.
_MANISKILL_DIR = lh.cache_dir() / "maniskill3"
os.environ.setdefault("MS_ASSET_DIR", str(_MANISKILL_DIR))
LEROBOT_SIM = lh.lerobot_sim("widowx", fps=5)


def task_names() -> list[str]:
    """The suite's task names, from LADDER's own registry."""
    from ladder.components.inference.robots.simpler.composition.specs import SPECS

    return sorted(spec.name for spec in SPECS)


class SimplerEnvLong(lh.LadderVectorEnv):
    suite = "simpler"
    robot = "widowx"
    resizable = False  # LADDER's SimplerEnv server renders its 480x640 camera
    episodes_per_task = EPISODES_PER_TASK
    fps = 5

    @classmethod
    def task_budget(cls, task: str) -> int:
        from ladder.components.inference.robots.simpler.composition.specs import SPECS

        return next(spec.max_steps for spec in SPECS if spec.name == task)

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._env = None
        # Once per process, before the server reports ready; the scene is built on the first reset.
        if os.environ["MS_ASSET_DIR"] == str(_MANISKILL_DIR):
            _stage_maniskill_assets()
        _seed_sapien_physx()
        _point_vulkan_at_nvidia()

    def _backend_reset(self, episode_ids: list[int]) -> list[Any]:
        if self._env is None:
            from ladder.components.inference.robots.simpler.env_builder import ABS_CONTROL_MODE
            from ladder.components.inference.robots.simpler.simpler_env import SimplerEnv, SimplerEnvConfig

            # The settings of LADDER's server config, robots/simpler/configs/widowx_composition.yaml.
            self._env = SimplerEnv(
                SimplerEnvConfig(
                    env_name=f"{self.task}-v0",
                    robot_name="widowx",
                    scene_name="bridge_table_1_v1",
                    rgb_overlay_path="",  # ManiSkill3 picks the overlay from env_name
                    sim_freq=500,
                    control_mode=ABS_CONTROL_MODE,
                    control_freq=5,
                    max_episode_steps=0,  # the task's own budget, 100 steps per subgoal
                    raytracing=False,
                    n_envs=self.num_envs,
                    episode_layouts=EPISODES_PER_TASK,
                )
            )
        return self._env.reset([e % EPISODES_PER_TASK for e in episode_ids])

    def _backend_step(self, actions: list[Any]) -> list[Any]:
        return self._env.step(actions)

    def _backend_close(self) -> None:
        if self._env is not None:
            self._env.close()
            self._env = None


def _stage_maniskill_assets() -> None:
    """ManiSkill reads one asset tree: LADDER's bridge scenes and ManiSkill's WidowX model side by side."""
    merged = _MANISKILL_DIR
    if not (merged / ".complete").exists():
        assets, widowx = lh.fetch("simpler_assets"), lh.fetch("widowx")
        (merged / "data" / "robots").mkdir(parents=True, exist_ok=True)
        for entry in (assets / "data").iterdir():
            link = merged / "data" / entry.name
            if entry.name != "robots":
                _link(link, entry)
        for entry in (assets / "data" / "robots").iterdir() if (assets / "data" / "robots").is_dir() else ():
            _link(merged / "data" / "robots" / entry.name, entry)
        _link(merged / "data" / "robots" / "widowx", widowx / "data" / "robots" / "widowx")
        (merged / ".complete").touch()


def _link(link, target) -> None:
    if not link.is_symlink():
        try:
            link.symlink_to(target)
        except FileExistsError:  # another process made it
            pass


def _seed_sapien_physx() -> None:
    """SAPIEN downloads its PhysX GPU library from GitHub on first use; use LADDER's pinned copy."""
    import sapien.physx

    library = os.path.join(os.path.expanduser("~"), ".sapien", "physx", sapien.physx.version(), "libPhysXGpu_64.so")
    if not os.path.exists(library):
        os.makedirs(os.path.dirname(library), exist_ok=True)
        temporary = f"{library}.{os.getpid()}"
        os.symlink(lh.fetch("sapien_physx") / "libPhysXGpu_64.so", temporary)
        os.replace(temporary, library)


def _point_vulkan_at_nvidia() -> None:
    """SAPIEN renders with Vulkan. A container may have the NVIDIA driver but no ICD file naming it."""
    if os.environ.get("VK_ICD_FILENAMES"):
        return
    for directory in ("/etc/vulkan/icd.d", "/usr/share/vulkan/icd.d"):
        if os.path.exists(os.path.join(directory, "nvidia_icd.json")):
            return
    icd = lh.cache_dir() / "nvidia_icd.json"
    if not icd.exists():
        icd.write_text('{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3.0"}}')
    os.environ["VK_ICD_FILENAMES"] = str(icd)


def vector_env_fns(
    task: str = "",
    n_envs: int = EPISODES_PER_TASK,
    **options: Any,
) -> dict[str, dict[int, Any]]:
    """One factory per task, ``{task: {0: factory}}``: every task, or the comma-separated ``task``.
    LeRobot builds only the tasks it runs (``--env.task``), so listing all 54 costs nothing.
    """
    names = lh.select_tasks(task, task_names())
    if os.environ["MS_ASSET_DIR"] == str(_MANISKILL_DIR):
        _stage_maniskill_assets()  # download before the server reports ready
    lh.fetch("sapien_physx")
    return {name: {0: partial(SimplerEnvLong, name, n_envs, **options)} for name in names}


def make_env(n_envs: int = EPISODES_PER_TASK, use_async_envs: bool = False, task: str = "", **options: Any):
    """One task as a batched vector env (``--env.use_native_batch``); LeRobot otherwise uses ``vector_env_fns``."""
    name = lh.one_task(task, task_names())
    return {name: {0: vector_env_fns(name, n_envs, **options)[name][0]()}}

	# Copyright (c) 2026 INSAIT. MIT License (see LICENSE).
	"""LADDER SimplerEnv-Long for LeRobot EnvHub: 54 long-horizon WidowX tasks on SAPIEN 3 / ManiSkill 3.

	Each task is one suite of LeRobot's ``{suite: {task_id: env}}``; its 24 episodes are the 24 object
	layouts LADDER defines (24 parallel envs on one GPU). See README.md.
	"""

	import os
	import re
	import sys

	REPO_ID = "INSAIT-Institute/LADDER-SimplerEnv-Long"

	_HERE = os.path.dirname(os.path.abspath(__file__))
	if not os.path.exists(os.path.join(_HERE, "_ladder_envhub.py")):
	# `make_env("<repo>")` downloads only this file: fetch the rest of the same revision.
	from huggingface_hub import snapshot_download

	_revision = os.path.basename(_HERE)
	_HERE = snapshot_download(REPO_ID, revision=_revision if re.fullmatch(r"[0-9a-f]{40}", _revision) else None)
	if _HERE not in sys.path:
	sys.path.insert(0, _HERE)

	from functools import partial # noqa: E402
	from typing import Any # noqa: E402

	import _ladder_envhub as lh # noqa: E402

	EPISODES_PER_TASK = 24
	# ManiSkill reads its asset directory once, when first imported; the tree is filled on the first reset.
	_MANISKILL_DIR = lh.cache_dir() / "maniskill3"
	os.environ.setdefault("MS_ASSET_DIR", str(_MANISKILL_DIR))
	LEROBOT_SIM = lh.lerobot_sim("widowx", fps=5)


	def task_names() -> list[str]:
	"""The suite's task names, from LADDER's own registry."""
	from ladder.components.inference.robots.simpler.composition.specs import SPECS

	return sorted(spec.name for spec in SPECS)


	class SimplerEnvLong(lh.LadderVectorEnv):
	suite = "simpler"
	robot = "widowx"
	resizable = False # LADDER's SimplerEnv server renders its 480x640 camera
	episodes_per_task = EPISODES_PER_TASK
	fps = 5

	@classmethod
	def task_budget(cls, task: str) -> int:
	from ladder.components.inference.robots.simpler.composition.specs import SPECS

	return next(spec.max_steps for spec in SPECS if spec.name == task)

	def __init__(self, args: Any, *kwargs: Any) -> None:
	super().__init__(args, *kwargs)
	self._env = None
	# Once per process, before the server reports ready; the scene is built on the first reset.
	if os.environ["MS_ASSET_DIR"] == str(_MANISKILL_DIR):
	_stage_maniskill_assets()
	_seed_sapien_physx()
	_point_vulkan_at_nvidia()

	def _backend_reset(self, episode_ids: list[int]) -> list[Any]:
	if self._env is None:
	from ladder.components.inference.robots.simpler.env_builder import ABS_CONTROL_MODE
	from ladder.components.inference.robots.simpler.simpler_env import SimplerEnv, SimplerEnvConfig

	# The settings of LADDER's server config, robots/simpler/configs/widowx_composition.yaml.
	self._env = SimplerEnv(
	SimplerEnvConfig(
	env_name=f"{self.task}-v0",
	robot_name="widowx",
	scene_name="bridge_table_1_v1",
	rgb_overlay_path="", # ManiSkill3 picks the overlay from env_name
	sim_freq=500,
	control_mode=ABS_CONTROL_MODE,
	control_freq=5,
	max_episode_steps=0, # the task's own budget, 100 steps per subgoal
	raytracing=False,
	n_envs=self.num_envs,
	episode_layouts=EPISODES_PER_TASK,
	)
	)
	return self._env.reset([e % EPISODES_PER_TASK for e in episode_ids])

	def _backend_step(self, actions: list[Any]) -> list[Any]:
	return self._env.step(actions)

	def _backend_close(self) -> None:
	if self._env is not None:
	self._env.close()
	self._env = None


	def _stage_maniskill_assets() -> None:
	"""ManiSkill reads one asset tree: LADDER's bridge scenes and ManiSkill's WidowX model side by side."""
	merged = _MANISKILL_DIR
	if not (merged / ".complete").exists():
	assets, widowx = lh.fetch("simpler_assets"), lh.fetch("widowx")
	(merged / "data" / "robots").mkdir(parents=True, exist_ok=True)
	for entry in (assets / "data").iterdir():
	link = merged / "data" / entry.name
	if entry.name != "robots":
	_link(link, entry)
	for entry in (assets / "data" / "robots").iterdir() if (assets / "data" / "robots").is_dir() else ():
	_link(merged / "data" / "robots" / entry.name, entry)
	_link(merged / "data" / "robots" / "widowx", widowx / "data" / "robots" / "widowx")
	(merged / ".complete").touch()


	def _link(link, target) -> None:
	if not link.is_symlink():
	try:
	link.symlink_to(target)
	except FileExistsError: # another process made it
	pass


	def _seed_sapien_physx() -> None:
	"""SAPIEN downloads its PhysX GPU library from GitHub on first use; use LADDER's pinned copy."""
	import sapien.physx

	library = os.path.join(os.path.expanduser("~"), ".sapien", "physx", sapien.physx.version(), "libPhysXGpu_64.so")
	if not os.path.exists(library):
	os.makedirs(os.path.dirname(library), exist_ok=True)
	temporary = f"{library}.{os.getpid()}"
	os.symlink(lh.fetch("sapien_physx") / "libPhysXGpu_64.so", temporary)
	os.replace(temporary, library)


	def _point_vulkan_at_nvidia() -> None:
	"""SAPIEN renders with Vulkan. A container may have the NVIDIA driver but no ICD file naming it."""
	if os.environ.get("VK_ICD_FILENAMES"):
	return
	for directory in ("/etc/vulkan/icd.d", "/usr/share/vulkan/icd.d"):
	if os.path.exists(os.path.join(directory, "nvidia_icd.json")):
	return
	icd = lh.cache_dir() / "nvidia_icd.json"
	if not icd.exists():
	icd.write_text('{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3.0"}}')
	os.environ["VK_ICD_FILENAMES"] = str(icd)


	def vector_env_fns(
	task: str = "",
	n_envs: int = EPISODES_PER_TASK,
	**options: Any,
	) -> dict[str, dict[int, Any]]:
	"""One factory per task, ``{task: {0: factory}}``: every task, or the comma-separated ``task``.

	LeRobot builds only the tasks it runs (``--env.task``), so listing all 54 costs nothing.
	"""
	names = lh.select_tasks(task, task_names())
	if os.environ["MS_ASSET_DIR"] == str(_MANISKILL_DIR):
	_stage_maniskill_assets() # download before the server reports ready
	lh.fetch("sapien_physx")
	return {name: {0: partial(SimplerEnvLong, name, n_envs, **options)} for name in names}


	def make_env(n_envs: int = EPISODES_PER_TASK, use_async_envs: bool = False, task: str = "", **options: Any):
	"""One task as a batched vector env (``--env.use_native_batch``); LeRobot otherwise uses ``vector_env_fns``."""
	name = lh.one_task(task, task_names())
	return {name: {0: vector_env_fns(name, n_envs, **options)[name][0]()}}