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

01171b7 verified 3 days ago

12.2 kB

# Copyright (c) 2026 INSAIT. MIT License (see LICENSE).
"""LADDER RoboLab-Long for LeRobot EnvHub: 44 long-horizon Franka (DROID) tasks on Isaac Sim 6 / Isaac Lab 3.
Each task is one suite of LeRobot's ``{suite: {task_id: env}}``; its 8 episodes run as 8 parallel envs
in one Isaac Lab world. Isaac Sim starts once per process. See README.md.
"""

import os
import re
import sys

REPO_ID = "INSAIT-Institute/LADDER-RoboLab-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)

import atexit  # noqa: E402
import logging  # noqa: E402
from concurrent.futures import ThreadPoolExecutor  # noqa: E402
from functools import partial  # noqa: E402
from typing import Any  # noqa: E402

import _ladder_envhub as lh  # noqa: E402

EPISODES_PER_TASK = 8
CONTROL_HZ = 15
LEROBOT_SIM = lh.lerobot_sim("franka", fps=CONTROL_HZ)


def _add_sources() -> None:
    """Isaac Lab, RoboLab (patched to read its assets from ROBOLAB_ASSET_DIR) and openpi_client, as LADDER pins them."""
    if _Isaac.sources:
        return
    robolab = lh.fetch("robolab_src")
    constants = robolab / "robolab" / "constants.py"
    text = constants.read_text()
    patched = 'ASSET_DIR = os.environ.get("ROBOLAB_ASSET_DIR") or os.path.join(PACKAGE_DIR, "assets")'
    if patched not in text:
        # LADDER's build_support/python/patches/robolab_constants_asset_dir.patch.
        text = text.replace('ASSET_DIR = os.path.join(PACKAGE_DIR, "assets")', patched, 1)
        text = text.replace(
            "    return os.path.join(PACKAGE_DIR, relative_path)",
            '    normalized = relative_path.replace(os.sep, "/")\n'
            '    if normalized.startswith("assets/"):\n'
            '        return os.path.join(ASSET_DIR, *normalized[len("assets/"):].split("/"))\n'
            "    return os.path.join(PACKAGE_DIR, relative_path)",
            1,
        )
        # Atomically: another process may be importing it.
        staged = constants.with_name(f".constants.{os.getpid()}.py")
        staged.write_text(text)
        os.replace(staged, constants)
    if patched not in constants.read_text():
        raise RuntimeError(f"{constants}: LADDER's asset-directory patch no longer applies")
    # The extracted wheel is an installed package on sys.path, with its dist-info (see _ladder_envhub.DOWNLOADS).
    lh.add_to_path(lh.fetch("isaaclab"), robolab, lh.fetch("openpi_src") / "packages" / "openpi-client" / "src")
    os.environ.setdefault("ROBOLAB_ASSET_DIR", str(lh.fetch("robolab_assets") / "assets"))
    _Isaac.sources = True


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

    return sorted(SPECS)


class _Isaac:
    """Isaac Sim starts once per process and runs on one thread; one task's world is kept at a time."""

    thread = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ladder-isaac")
    sources = False
    app: Any = None
    env: Any = None
    env_cfg: Any = None
    oracle: Any = None
    key: Any = None


def _start_isaac() -> None:
    if _Isaac.app is not None:
        return
    _add_sources()
    # Running this environment accepts the NVIDIA Isaac Sim EULA (see README); Kit asks on stdin otherwise.
    os.environ.setdefault("OMNI_KIT_ACCEPT_EULA", "YES")
    os.environ.setdefault("ACCEPT_EULA", "Y")
    import cv2  # noqa: F401
    from isaaclab.app import AppLauncher

    root = logging.getLogger()
    levels = {handler: handler.level for handler in root.handlers}
    with lh.ignore_signal_in_thread():
        _Isaac.app = AppLauncher(headless=True, enable_cameras=True).app
    for handler, level in levels.items():  # Kit lowers the levels of existing handlers
        handler.setLevel(level)
    # Kit hangs in its own shutdown; leave without it, keeping the exit status.
    atexit.register(
        lambda: (
            sys.stdout.flush(),
            sys.stderr.flush(),
            os._exit(0 if getattr(sys, "last_value", None) is None else 1),
        )
    )
    import carb.settings
    import robolab.constants

    # As LADDER's env_server.py: LADDER scores subgoals itself, and renderer failures stay visible.
    robolab.constants.ENABLE_SUBTASK_PROGRESS_CHECKING = False
    carb.settings.get_settings().set("/persistent/app/usd/muteUsdDiagnostics", False)


class RoboLabLong(lh.LadderVectorEnv):
    suite = "robolab"
    robot = "franka"
    default_image_size = (224, 224)
    episodes_per_task = EPISODES_PER_TASK
    fps = CONTROL_HZ

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

        return int(SPECS[task].episode_length_s * CONTROL_HZ)

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._env = None
        self._raw = None
        self._converter = None
        # LeRobot's server builds an env (one env of the first task) before it reports ready. Isaac Sim starts here,
        # and the first env of a process builds its world here too: that compiles the shaders and loads the assets
        # every task shares, which on a new machine or container would otherwise fall in a client's first open.
        # A client's own batch size gets its world on its first reset.
        self._fresh = _Isaac.thread.submit(self._start).result()

    def _start(self) -> bool:
        _start_isaac()
        return _Isaac.env is None and self._build()

    def goal_sources(self) -> list[str | None] | None:
        # The object each goal places (a relay's legs are separate goals), as LADDER's privileged.py names it.
        return [
            next((geometry[0] for _, geometry in reversed(list(goal.legs)) if geometry is not None), None)
            for goal in _Isaac.oracle or []
        ]

    def _backend_reset(self, episode_ids: list[int]) -> list[Any]:
        return _Isaac.thread.submit(self._reset).result()

    def _backend_step(self, actions: list[Any]) -> list[Any]:
        if _Isaac.env is not self._env:
            raise RuntimeError("Another task's world replaced this one: run one task at a time (--env.max_parallel_tasks=1)")
        return _Isaac.thread.submit(self._step, actions).result()

    def _build(self) -> bool:
        """Build this task's world at this batch size, unless the process holds it; True if it was built."""
        from ladder.components.inference.robots.robolab.composition.specs import task_class_name, task_target
        from ladder.components.inference.robots.robolab.oracle import build_oracle
        from robolab.core.environments.runtime import create_env
        from robolab.registrations.droid.auto_env_registrations_jointpos import auto_register_droid_envs

        self._class_name = task_class_name(self.task)
        key = (self._class_name, self.num_envs)
        if _Isaac.key == key:
            return False
        if _Isaac.env is not None:
            _Isaac.env.close()
            _Isaac.env = _Isaac.key = None
        auto_register_droid_envs(task_dirs=None, task=[task_target(self.task)], cameras=None)
        # The defaults of LADDER's env_server.py (RoboLab's add_common_eval_args).
        _Isaac.env, _Isaac.env_cfg = create_env(
            self._class_name,
            device="cuda:0",
            num_envs=self.num_envs,
            instruction_type="default",
            policy="ladder_runner",
            renderer="realtime",
            rendering_mode="balanced",
        )
        _Isaac.oracle = build_oracle(getattr(_Isaac.env.cfg, "subtasks", None), _Isaac.env_cfg.instruction)
        _Isaac.key = key
        if int(_Isaac.env.max_episode_length) != self._max_episode_steps:
            logging.getLogger(__name__).warning(
                "%s: Isaac Lab budget %d steps, expected %d", self.task, _Isaac.env.max_episode_length,
                self._max_episode_steps,
            )
        return True

    def _reset(self) -> list[Any]:
        from ladder.components.inference.robots.robolab.droid_conversions import DroidJointposConverter

        # As LADDER's env_server.py: a new world is reset as it is, a kept one first unfreezes its finished envs.
        if not self._build() and not self._fresh:
            _Isaac.env.reset_eval_state()
        self._fresh = False
        self._env = _Isaac.env
        self._env.reset()
        self._raw, _ = self._env.reset()  # RoboLab's own double reset (robolab.eval.episode.run_episode)
        # DroidJointposConverter takes (width, height), as LADDER's --image-size.
        self._converter = DroidJointposConverter(num_envs=self.num_envs, image_size=(self._w, self._h))
        return self._observations()

    def _step(self, actions: list[Any]) -> list[Any]:
        import omni.kit.app
        import omni.timeline
        import torch

        joints = self._raw["proprio_obs"]["arm_joint_pos"].detach().cpu().numpy()
        commands = self._converter.actions_to_jointpos(actions, joints)
        timeline, kit = omni.timeline.get_timeline_interface(), omni.kit.app.get_app()
        while not timeline.is_playing():
            kit.update()
        self._raw, _, _, _, _ = self._env.step(torch.as_tensor(commands, device=self._env.device))
        return self._observations()

    def _observations(self) -> list[Any]:
        images, proprio = self._raw["image_obs"], self._raw["proprio_obs"]
        results = self._env.get_env_results()
        active = set(self._env.active_env_ids)
        # During a step, _step_count is not yet incremented; LADDER's server counts the step first.
        out_of_time = self._step_count + 1 >= int(self._env.max_episode_length)
        achieved = None
        if _Isaac.oracle:
            achieved = [[bool(g.check(self._env, env_id=i)) for g in _Isaac.oracle] for i in range(self.num_envs)]
        observations = []
        for i in range(self.num_envs):
            success = bool(results[i]["success"])
            metadata = {
                "env_name": self._class_name,
                "robot_name": "franka",
                "instruction": _Isaac.env_cfg.instruction,
                "done": success,
                "truncated": out_of_time or i not in active,
                "oracle": {"achieved": achieved[i] if achieved is not None else [success]},
            }
            observations.append(
                self._converter.observation_dict(
                    env_i=i,
                    main_image=images["over_shoulder_left_camera"][i].detach().cpu().numpy(),
                    wrist_image=images["wrist_cam"][i].detach().cpu().numpy(),
                    joints=proprio["arm_joint_pos"][i].detach().cpu().numpy(),
                    gripper_position=proprio["gripper_pos"][i].detach().cpu().numpy(),
                    metadata=metadata,
                )
            )
        return observations


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 44 costs nothing.
    """
    names = lh.select_tasks(task, task_names())
    _add_sources()  # download before the server reports ready
    return {name: {0: partial(RoboLabLong, 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 RoboLab-Long for LeRobot EnvHub: 44 long-horizon Franka (DROID) tasks on Isaac Sim 6 / Isaac Lab 3.

	Each task is one suite of LeRobot's ``{suite: {task_id: env}}``; its 8 episodes run as 8 parallel envs
	in one Isaac Lab world. Isaac Sim starts once per process. See README.md.
	"""

	import os
	import re
	import sys

	REPO_ID = "INSAIT-Institute/LADDER-RoboLab-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)

	import atexit # noqa: E402
	import logging # noqa: E402
	from concurrent.futures import ThreadPoolExecutor # noqa: E402
	from functools import partial # noqa: E402
	from typing import Any # noqa: E402

	import _ladder_envhub as lh # noqa: E402

	EPISODES_PER_TASK = 8
	CONTROL_HZ = 15
	LEROBOT_SIM = lh.lerobot_sim("franka", fps=CONTROL_HZ)


	def _add_sources() -> None:
	"""Isaac Lab, RoboLab (patched to read its assets from ROBOLAB_ASSET_DIR) and openpi_client, as LADDER pins them."""
	if _Isaac.sources:
	return
	robolab = lh.fetch("robolab_src")
	constants = robolab / "robolab" / "constants.py"
	text = constants.read_text()
	patched = 'ASSET_DIR = os.environ.get("ROBOLAB_ASSET_DIR") or os.path.join(PACKAGE_DIR, "assets")'
	if patched not in text:
	# LADDER's build_support/python/patches/robolab_constants_asset_dir.patch.
	text = text.replace('ASSET_DIR = os.path.join(PACKAGE_DIR, "assets")', patched, 1)
	text = text.replace(
	" return os.path.join(PACKAGE_DIR, relative_path)",
	' normalized = relative_path.replace(os.sep, "/")\n'
	' if normalized.startswith("assets/"):\n'
	' return os.path.join(ASSET_DIR, *normalized[len("assets/"):].split("/"))\n'
	" return os.path.join(PACKAGE_DIR, relative_path)",
	1,
	)
	# Atomically: another process may be importing it.
	staged = constants.with_name(f".constants.{os.getpid()}.py")
	staged.write_text(text)
	os.replace(staged, constants)
	if patched not in constants.read_text():
	raise RuntimeError(f"{constants}: LADDER's asset-directory patch no longer applies")
	# The extracted wheel is an installed package on sys.path, with its dist-info (see _ladder_envhub.DOWNLOADS).
	lh.add_to_path(lh.fetch("isaaclab"), robolab, lh.fetch("openpi_src") / "packages" / "openpi-client" / "src")
	os.environ.setdefault("ROBOLAB_ASSET_DIR", str(lh.fetch("robolab_assets") / "assets"))
	_Isaac.sources = True


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

	return sorted(SPECS)


	class _Isaac:
	"""Isaac Sim starts once per process and runs on one thread; one task's world is kept at a time."""

	thread = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ladder-isaac")
	sources = False
	app: Any = None
	env: Any = None
	env_cfg: Any = None
	oracle: Any = None
	key: Any = None


	def _start_isaac() -> None:
	if _Isaac.app is not None:
	return
	_add_sources()
	# Running this environment accepts the NVIDIA Isaac Sim EULA (see README); Kit asks on stdin otherwise.
	os.environ.setdefault("OMNI_KIT_ACCEPT_EULA", "YES")
	os.environ.setdefault("ACCEPT_EULA", "Y")
	import cv2 # noqa: F401
	from isaaclab.app import AppLauncher

	root = logging.getLogger()
	levels = {handler: handler.level for handler in root.handlers}
	with lh.ignore_signal_in_thread():
	_Isaac.app = AppLauncher(headless=True, enable_cameras=True).app
	for handler, level in levels.items(): # Kit lowers the levels of existing handlers
	handler.setLevel(level)
	# Kit hangs in its own shutdown; leave without it, keeping the exit status.
	atexit.register(
	lambda: (
	sys.stdout.flush(),
	sys.stderr.flush(),
	os._exit(0 if getattr(sys, "last_value", None) is None else 1),
	)
	)
	import carb.settings
	import robolab.constants

	# As LADDER's env_server.py: LADDER scores subgoals itself, and renderer failures stay visible.
	robolab.constants.ENABLE_SUBTASK_PROGRESS_CHECKING = False
	carb.settings.get_settings().set("/persistent/app/usd/muteUsdDiagnostics", False)


	class RoboLabLong(lh.LadderVectorEnv):
	suite = "robolab"
	robot = "franka"
	default_image_size = (224, 224)
	episodes_per_task = EPISODES_PER_TASK
	fps = CONTROL_HZ

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

	return int(SPECS[task].episode_length_s * CONTROL_HZ)

	def __init__(self, args: Any, *kwargs: Any) -> None:
	super().__init__(args, *kwargs)
	self._env = None
	self._raw = None
	self._converter = None
	# LeRobot's server builds an env (one env of the first task) before it reports ready. Isaac Sim starts here,
	# and the first env of a process builds its world here too: that compiles the shaders and loads the assets
	# every task shares, which on a new machine or container would otherwise fall in a client's first open.
	# A client's own batch size gets its world on its first reset.
	self._fresh = _Isaac.thread.submit(self._start).result()

	def _start(self) -> bool:
	_start_isaac()
	return _Isaac.env is None and self._build()

	def goal_sources(self) -> list[str \| None] \| None:
	# The object each goal places (a relay's legs are separate goals), as LADDER's privileged.py names it.
	return [
	next((geometry[0] for _, geometry in reversed(list(goal.legs)) if geometry is not None), None)
	for goal in _Isaac.oracle or []
	]

	def _backend_reset(self, episode_ids: list[int]) -> list[Any]:
	return _Isaac.thread.submit(self._reset).result()

	def _backend_step(self, actions: list[Any]) -> list[Any]:
	if _Isaac.env is not self._env:
	raise RuntimeError("Another task's world replaced this one: run one task at a time (--env.max_parallel_tasks=1)")
	return _Isaac.thread.submit(self._step, actions).result()

	def _build(self) -> bool:
	"""Build this task's world at this batch size, unless the process holds it; True if it was built."""
	from ladder.components.inference.robots.robolab.composition.specs import task_class_name, task_target
	from ladder.components.inference.robots.robolab.oracle import build_oracle
	from robolab.core.environments.runtime import create_env
	from robolab.registrations.droid.auto_env_registrations_jointpos import auto_register_droid_envs

	self._class_name = task_class_name(self.task)
	key = (self._class_name, self.num_envs)
	if _Isaac.key == key:
	return False
	if _Isaac.env is not None:
	_Isaac.env.close()
	_Isaac.env = _Isaac.key = None
	auto_register_droid_envs(task_dirs=None, task=[task_target(self.task)], cameras=None)
	# The defaults of LADDER's env_server.py (RoboLab's add_common_eval_args).
	_Isaac.env, _Isaac.env_cfg = create_env(
	self._class_name,
	device="cuda:0",
	num_envs=self.num_envs,
	instruction_type="default",
	policy="ladder_runner",
	renderer="realtime",
	rendering_mode="balanced",
	)
	_Isaac.oracle = build_oracle(getattr(_Isaac.env.cfg, "subtasks", None), _Isaac.env_cfg.instruction)
	_Isaac.key = key
	if int(_Isaac.env.max_episode_length) != self._max_episode_steps:
	logging.getLogger(__name__).warning(
	"%s: Isaac Lab budget %d steps, expected %d", self.task, _Isaac.env.max_episode_length,
	self._max_episode_steps,
	)
	return True

	def _reset(self) -> list[Any]:
	from ladder.components.inference.robots.robolab.droid_conversions import DroidJointposConverter

	# As LADDER's env_server.py: a new world is reset as it is, a kept one first unfreezes its finished envs.
	if not self._build() and not self._fresh:
	_Isaac.env.reset_eval_state()
	self._fresh = False
	self._env = _Isaac.env
	self._env.reset()
	self._raw, _ = self._env.reset() # RoboLab's own double reset (robolab.eval.episode.run_episode)
	# DroidJointposConverter takes (width, height), as LADDER's --image-size.
	self._converter = DroidJointposConverter(num_envs=self.num_envs, image_size=(self._w, self._h))
	return self._observations()

	def _step(self, actions: list[Any]) -> list[Any]:
	import omni.kit.app
	import omni.timeline
	import torch

	joints = self._raw["proprio_obs"]["arm_joint_pos"].detach().cpu().numpy()
	commands = self._converter.actions_to_jointpos(actions, joints)
	timeline, kit = omni.timeline.get_timeline_interface(), omni.kit.app.get_app()
	while not timeline.is_playing():
	kit.update()
	self._raw, _, _, _, _ = self._env.step(torch.as_tensor(commands, device=self._env.device))
	return self._observations()

	def _observations(self) -> list[Any]:
	images, proprio = self._raw["image_obs"], self._raw["proprio_obs"]
	results = self._env.get_env_results()
	active = set(self._env.active_env_ids)
	# During a step, _step_count is not yet incremented; LADDER's server counts the step first.
	out_of_time = self._step_count + 1 >= int(self._env.max_episode_length)
	achieved = None
	if _Isaac.oracle:
	achieved = [[bool(g.check(self._env, env_id=i)) for g in _Isaac.oracle] for i in range(self.num_envs)]
	observations = []
	for i in range(self.num_envs):
	success = bool(results[i]["success"])
	metadata = {
	"env_name": self._class_name,
	"robot_name": "franka",
	"instruction": _Isaac.env_cfg.instruction,
	"done": success,
	"truncated": out_of_time or i not in active,
	"oracle": {"achieved": achieved[i] if achieved is not None else [success]},
	}
	observations.append(
	self._converter.observation_dict(
	env_i=i,
	main_image=images["over_shoulder_left_camera"][i].detach().cpu().numpy(),
	wrist_image=images["wrist_cam"][i].detach().cpu().numpy(),
	joints=proprio["arm_joint_pos"][i].detach().cpu().numpy(),
	gripper_position=proprio["gripper_pos"][i].detach().cpu().numpy(),
	metadata=metadata,
	)
	)
	return observations


	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 44 costs nothing.
	"""
	names = lh.select_tasks(task, task_names())
	_add_sources() # download before the server reports ready
	return {name: {0: partial(RoboLabLong, 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]()}}