File size: 5,510 Bytes
73123f2 fcf6ade 73123f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""Controller capabilities the viewer drives that lerobot does not advertise.
Everything here works against an unmodified ``lerobot`` from main. The demo
pins upstream rather than carrying patches into it, so anything the UI needs
beyond the ``RobotController`` protocol is adapted from the outside.
"""
from __future__ import annotations
import ast
import importlib
import logging
from pathlib import Path
from lerobot.robots.unitree_g1 import unitree_g1
logger = logging.getLogger(__name__)
# Scalar commands a controller keeps on an instance attribute, as
# {controller class: {action key: (min, max, default, attribute)}}.
#
# The protocol has no vocabulary for these: run_step is handed the joystick
# axes, and anything else the policy takes as an observation is reachable only
# as an attribute. The driver applies them by assignment, hence the name.
COMMANDS: dict[str, dict[str, tuple[float, float, float, str]]] = {
# Commanded base height, i.e. how deep GR00T squats. The policy reads it as
# an observation and holds whatever it is given rather than tracking to it,
# so a slider maps straight onto it. Upstream only ever nudges it by
# 0.001/tick from held gamepad buttons, which no browser is going to send.
"GrootLocomotionController": {"command.height": (0.50, 1.00, 0.74, "groot_height_cmd")},
}
def limit_inference_threads(threads: int = 1) -> None:
"""Pin ONNX Runtime to ``threads`` per pool, for every session made later.
Thread pools are sized from the machine's core count, and inside a container
that is the host's count rather than the CPU quota. On a 2-vCPU Space that
meant ~60 threads per process contending for 2 cores, which costs far more
than the inference itself: physics needs about 5% of a core here, yet the
container sat pegged at its limit and the policies missed their deadlines.
These are small MLPs, so one thread each beats any amount of contention.
lerobot builds its sessions as ``ort.InferenceSession(path)`` with no
options, so the only way in from outside is to wrap the constructor.
The matching BLAS and OpenMP limits have to be environment variables set
before numpy and torch are imported, so they live in the Dockerfile.
"""
import onnxruntime as ort
if getattr(ort.InferenceSession, "_thread_capped", False):
return
original = ort.InferenceSession
def capped(*args, **kwargs):
# sess_options is the second positional parameter; only inject when the
# caller passed none, so an explicit choice is never overridden.
if len(args) < 2 and kwargs.get("sess_options") is None:
options = ort.SessionOptions()
options.intra_op_num_threads = threads
options.inter_op_num_threads = threads
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
kwargs["sess_options"] = options
return original(*args, **kwargs)
capped._thread_capped = True
ort.InferenceSession = capped
logger.info(f"onnxruntime capped to {threads} thread(s) per session")
def _base_names(node: ast.ClassDef) -> set[str]:
"""Base classes of a ClassDef as written (``X`` or ``mod.X`` -> ``X``)."""
return {b.id for b in node.bases if isinstance(b, ast.Name)} | {
b.attr for b in node.bases if isinstance(b, ast.Attribute)
}
def discover_controllers() -> dict[str, str]:
"""Map controller class name to module path, from the installed lerobot.
Anything in ``controllers/`` declaring a ``RobotController`` subclass is
selectable, so dropping a module in there is enough; nothing here needs
editing. Deliberately uncached, so a new one appears on a page refresh
rather than a restart.
The modules are parsed rather than imported: ``controllers/__init__.py``
imports all of them eagerly, which pulls in onnxruntime and every policy's
dependencies just to answer "what can I pick?", and would let a single
controller with an unmet optional dependency hide all the others.
"""
package = f"{unitree_g1.__package__}.controllers"
found: dict[str, str] = {}
for path in sorted((Path(unitree_g1.__file__).parent / "controllers").glob("[!_]*.py")):
try:
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
except (OSError, SyntaxError) as e:
logger.warning(f"skipping controller module {path.name}: {e}")
continue
for n in tree.body:
if isinstance(n, ast.ClassDef) and "RobotController" in _base_names(n):
found[n.name] = f"{package}.{path.stem}"
return found
def available_controllers() -> list[str]:
"""Controller class names that can be selected."""
return sorted(discover_controllers())
def make_controller(name: str):
"""Instantiate a discovered controller by class name.
Mirrors lerobot's ``make_robot_controller``, which is not reused because it
matches against a hardcoded list of the three controllers that shipped with
it: a newly added one would be discovered and offered in the UI, then be
rejected on click. Going through the discovered module path means adding a
controller really is just dropping in a file.
"""
modules = discover_controllers()
if name not in modules:
raise ValueError(f"Unknown controller: {name!r}. Available: {sorted(modules)}")
return getattr(importlib.import_module(modules[name]), name)()
|