Pass the visualizer windows to the script input helpers explicitly

The teleop session now reads keyboard input only from the visualizers it
is given. open_teleop_device takes them and opens the device through
select_teleop_cfg and create_teleop_device, and add_key_callbacks takes
them instead of looking up the simulation. The replay, annotate,
consolidated Mimic and pick-and-place scripts pass sim.visualizers.

Author: rwiltz

examples/demos/pick_and_place.py

@@ -333,9 +333,10 @@ def main() -> None:
             # the keyboard runs through Isaac Capture, which the teleop extra installs
             from isaaclab_teleop.script_input import open_teleop_device
 
-            teleop = teleop_device.enter_context(open_teleop_device(teleop_cfg))
+            teleop = teleop_device.enter_context(open_teleop_device(teleop_cfg, pick_and_place.sim.visualizers))
             add_key_callbacks(
-                {"N": lambda: pick_and_place.auto_aim(cube=True), "M": lambda: pick_and_place.auto_aim(cube=False)}
+                pick_and_place.sim.visualizers,
+                {"N": lambda: pick_and_place.auto_aim(cube=True), "M": lambda: pick_and_place.auto_aim(cube=False)},
             )
             print(teleop)
             print("Pick up the purple cube and drop it on the red sphere, in ALL environments at once.")

scripts/imitation_learning/isaaclab_mimic/annotate_demos.py

@@ -288,7 +288,7 @@ def annotate_dataset(env_cfg, dataset_file_handler: HDF5DatasetFileHandler, succ
         key_callbacks = {"N": play_cb, "B": pause_cb, "Q": skip_episode_cb}
         if not args_cli.auto:
             key_callbacks["S"] = mark_subtask_cb
-        add_key_callbacks(key_callbacks)
+        add_key_callbacks(env.unwrapped.sim.visualizers, key_callbacks)
 
     # simulate environment -- run everything in inference mode
     exported_episode_count = 0

scripts/imitation_learning/isaaclab_mimic/consolidated_demo.py

@@ -458,7 +458,9 @@ def run_consolidated_demo(env_cfg, success_term):
     if args_cli.teleop_env_index is not None:
         from isaaclab_teleop.script_input import open_teleop_device
 
-        teleop_interface = teleop_device.enter_context(open_teleop_device(teleop_device_cfg()))
+        teleop_interface = teleop_device.enter_context(
+            open_teleop_device(teleop_device_cfg(), env.unwrapped.sim.visualizers)
+        )
 
     # make data generator object
     data_generator_asyncio_tasks = []

scripts/tools/replay_demos.py

@@ -311,7 +311,7 @@ def replay_dataset(
     # create environment from loaded config
     env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
 
-    add_key_callbacks({"N": play_cb, "B": pause_cb})
+    add_key_callbacks(env.unwrapped.sim.visualizers, {"N": play_cb, "B": pause_cb})
     print('Press "B" to pause and "N" to resume the replayed actions.')
 
     # Determine if state validation should be conducted

source/isaaclab/isaaclab/visualizers/key_event_source.py

@@ -8,7 +8,11 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Mapping, Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .base_visualizer import BaseVisualizer
 
 logger = logging.getLogger(__name__)
 
@@ -121,24 +125,24 @@ def _emit_focus_lost(self) -> None:
                 _call(subscription._on_focus_lost)
 
 
-def add_key_callbacks(callbacks: Mapping[str, Callable[[], None]]) -> Callable[[], None]:
+def add_key_callbacks(
+    visualizers: Sequence[BaseVisualizer], callbacks: Mapping[str, Callable[[], None]]
+) -> Callable[[], None]:
     """Call ``callbacks[key]`` when ``key`` is pressed in a visualizer window, e.g. for a script's hotkeys.
 
     Keys are letters or digits (``"N"``, ``"1"``) or W3C ``KeyboardEvent.code`` strings (``"Escape"``).
     Key repeat is not a press, and a press reported by several visualizers (Kit visualizers share the
-    app window) fires once. Reads every visualizer window open when called; the callbacks end when
-    the visualizers close, or when the returned function is called.
+    app window) fires once. The callbacks end when the visualizers close, or when the returned
+    function is called.
 
     Args:
+        visualizers: The windows to read, usually ``sim.visualizers``.
         callbacks: The function to call for each key. Takes no arguments.
 
     Returns:
         A function that removes the callbacks.
     """
-    from isaaclab.sim import SimulationContext
-
     by_code = {_w3c_code(key): func for key, func in callbacks.items()}
-    sim = SimulationContext.instance()
     # shared by every window, so a press several visualizers report fires once
     held: set[str] = set()
 
@@ -151,7 +155,7 @@ def on_key(code: str, pressed: bool) -> None:
                 by_code[code]()
 
     subscriptions = []
-    for visualizer in sim.visualizers if sim is not None else ():
+    for visualizer in visualizers:
         source = visualizer.key_event_source
         if source is not None:
             subscriptions.append(source.add_key_listener(on_key, held.clear))

source/isaaclab/test/visualizers/test_key_event_source.py

@@ -12,7 +12,6 @@
 
 import logging
 from types import SimpleNamespace
-from unittest.mock import patch
 
 import pytest
 
@@ -149,13 +148,15 @@ def source():
     sim = SimpleNamespace(
         visualizers=[SimpleNamespace(key_event_source=key_event_source), SimpleNamespace(key_event_source=None)]
     )
-    with patch("isaaclab.sim.SimulationContext.instance", return_value=sim):
-        yield key_event_source
+    return key_event_source, sim.visualizers
 
 
 def test_key_callbacks_fire_once_per_press(source):
     pressed = []
-    add_key_callbacks({"N": lambda: pressed.append("N"), "1": lambda: pressed.append("1"), "Escape": lambda: 0})
+    source, visualizers = source
+    add_key_callbacks(
+        visualizers, {"N": lambda: pressed.append("N"), "1": lambda: pressed.append("1"), "Escape": lambda: 0}
+    )
 
     source._emit_key("KeyN", True)
     source._emit_key("KeyN", True)  # key repeat
@@ -168,15 +169,17 @@ def test_key_callbacks_fire_once_per_press(source):
 
 def test_focus_loss_ends_a_held_key(source):
     pressed = []
-    add_key_callbacks({"N": lambda: pressed.append("N")})
+    source, visualizers = source
+    add_key_callbacks(visualizers, {"N": lambda: pressed.append("N")})
     source._emit_key("KeyN", True)
     source._emit_focus_lost()
     source._emit_key("KeyN", True)
     assert pressed == ["N", "N"]
 
 
 def test_removing_the_key_callbacks_unsubscribes(source):
-    remove = add_key_callbacks({"N": lambda: None})
+    source, visualizers = source
+    remove = add_key_callbacks(visualizers, {"N": lambda: None})
     assert len(source._subscriptions) == 1
     remove()
     assert source._subscriptions == []

source/isaaclab_teleop/isaaclab_teleop/script_input.py

@@ -8,17 +8,19 @@
 from __future__ import annotations
 
 import contextlib
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from typing import TYPE_CHECKING, Any
 
-from .device_selection import _require_isaaccapture, pipeline_from_device_cfg
+from .device_selection import create_teleop_device, select_teleop_cfg
+from .isaac_teleop_cfg import CLOUDXR_STANDALONE_ENV, IsaacTeleopCfg
 
 if TYPE_CHECKING:
     from isaaclab.devices import DeviceCfg
+    from isaaclab.visualizers import BaseVisualizer
 
 
 @contextlib.contextmanager
-def open_teleop_device(device_cfg: DeviceCfg) -> Iterator[Any]:
+def open_teleop_device(device_cfg: DeviceCfg, visualizers: Sequence[BaseVisualizer]) -> Iterator[Any]:
     """Open the keyboard, gamepad or SpaceMouse a device config describes, through Isaac Capture.
 
     The device runs as an Isaac Capture pipeline, started without a headset on the standalone CloudXR
@@ -28,6 +30,7 @@ def open_teleop_device(device_cfg: DeviceCfg) -> Iterator[Any]:
 
     Args:
         device_cfg: A keyboard, gamepad or SpaceMouse config, e.g. :class:`~isaaclab.devices.Se3KeyboardCfg`.
+        visualizers: The windows that supply keyboard input, usually ``sim.visualizers``.
 
     Yields:
         The open device.
@@ -36,15 +39,11 @@ def open_teleop_device(device_cfg: DeviceCfg) -> Iterator[Any]:
         ImportError: If Isaac Capture is not installed.
         ValueError: If no Isaac Capture pipeline serves the device config.
     """
-    _require_isaaccapture()
-    pipeline = pipeline_from_device_cfg(device_cfg, device_cfg.sim_device)
-    if pipeline is None:
+    cfg = select_teleop_cfg(None, "device", {"device": device_cfg}, device_cfg.sim_device)
+    if not isinstance(cfg, IsaacTeleopCfg):
         raise ValueError(f"No Isaac Capture pipeline serves {type(device_cfg).__name__}")
-    from .isaac_teleop_cfg import CLOUDXR_STANDALONE_ENV
-    from .isaac_teleop_device import create_isaac_teleop_device
-
-    with create_isaac_teleop_device(
-        pipeline, use_kit_xr_bridge=False, cloudxr_env_file=CLOUDXR_STANDALONE_ENV
+    with create_teleop_device(
+        cfg, use_kit_xr_bridge=False, cloudxr_env_file=CLOUDXR_STANDALONE_ENV, visualizers=visualizers
     ) as device:
         device.request_start()
         yield device

source/isaaclab_teleop/test/test_script_input.py

@@ -9,36 +9,39 @@
 
 from __future__ import annotations
 
-from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 import isaaclab_teleop.device_selection as device_selection
 import isaaclab_teleop.script_input as script_input
 import pytest
-from isaaclab_teleop.isaac_teleop_cfg import CLOUDXR_STANDALONE_ENV
+from isaaclab_teleop.isaac_teleop_cfg import CLOUDXR_STANDALONE_ENV, IsaacTeleopCfg
 from isaaclab_teleop.script_input import open_teleop_device
 
+from isaaclab.devices import Se3KeyboardCfg
+
 pytestmark = pytest.mark.unit
 
 
 def test_without_isaacteleop_opening_a_device_raises(monkeypatch):
-    monkeypatch.setattr(device_selection, "_isaacteleop_available", lambda: False)
+    monkeypatch.setattr(device_selection.importlib.util, "find_spec", lambda _: None)
     with pytest.raises(ImportError, match="teleop extra"):
-        with open_teleop_device(SimpleNamespace(name="keyboard", sim_device="cpu")):
+        with open_teleop_device(Se3KeyboardCfg(), visualizers=()):
             pass
 
 
 def test_with_isaacteleop_the_pipeline_runs_headset_free(monkeypatch):
-    monkeypatch.setattr(device_selection, "_isaacteleop_available", lambda: True)
-    monkeypatch.setattr(script_input, "pipeline_from_device_cfg", lambda cfg, sim_device: f"pipeline:{cfg.name}")
+    cfg = IsaacTeleopCfg()
+    monkeypatch.setattr(script_input, "select_teleop_cfg", lambda env_cfg, name, defaults, sim_device: cfg)
     device = MagicMock()
     device.__enter__.return_value = device
-    with patch("isaaclab_teleop.isaac_teleop_device.create_isaac_teleop_device", return_value=device) as create:
-        with open_teleop_device(SimpleNamespace(name="keyboard", sim_device="cpu")) as opened:
-            assert opened is device
-            device.request_start.assert_called_once()
-            device.__exit__.assert_not_called()
+    create = MagicMock(return_value=device)
+    monkeypatch.setattr(script_input, "create_teleop_device", create)
+    visualizers = [object()]
+    with open_teleop_device(Se3KeyboardCfg(), visualizers) as opened:
+        assert opened is device
+        device.request_start.assert_called_once()
+        device.__exit__.assert_not_called()
     create.assert_called_once_with(
-        "pipeline:keyboard", use_kit_xr_bridge=False, cloudxr_env_file=CLOUDXR_STANDALONE_ENV
+        cfg, use_kit_xr_bridge=False, cloudxr_env_file=CLOUDXR_STANDALONE_ENV, visualizers=visualizers
     )
     device.__exit__.assert_called_once()