diff --git a/docs/INSTALL.md b/docs/INSTALL.md index ec7c0dc..f7dea5b 100644 --- a/docs/INSTALL.md +++ b/docs/INSTALL.md @@ -11,6 +11,7 @@ | ๐Ÿงช IsaacLab / Isaac Sim | [Install Isaac Sim / IsaacLab](#isaaclab--isaac-sim) | | ๐Ÿฆพ Franka real | Set up the [realtime kernel](https://frankarobotics.github.io/docs/libfranka/docs/real_time_kernel.html), then install `mpail2` with `".[franka]"` | | ๐Ÿค– Kinova real | Install [ROS 2](https://docs.ros.org/en/humble/Installation/Ubuntu-Install-Debs.html) and [`ros2_kortex`](https://github.com/Kinovarobotics/ros2_kortex) | +| ๐Ÿฆฟ SO-101 real | Install `mpail2` with `".[so101]"`, plus a separate [`lerobot`](https://github.com/huggingface/lerobot) clone โ€” see [`mpail2/envs/real/so101/README.md`](../mpail2/envs/real/so101/README.md) | ## Default Install @@ -31,6 +32,7 @@ IsaacLab and the real-robot environments require extra setup: - IsaacLab / Isaac Sim: install Isaac Sim / IsaacLab first, then install `mpail2` - Franka real robot: install `mpail2` with `".[franka]"` - Kinova real robot: install ROS 2 / `ros2_kortex` +- SO-101 real robot: install `mpail2` with `".[so101]"` โ€” see [`mpail2/envs/real/so101/README.md`](../mpail2/envs/real/so101/README.md) for the full setup (needs a separate `lerobot` clone + conda env) ## Gym / MuJoCo diff --git a/mpail2/configs/cfgs.py b/mpail2/configs/cfgs.py index 190a5ca..17f5039 100644 --- a/mpail2/configs/cfgs.py +++ b/mpail2/configs/cfgs.py @@ -189,6 +189,13 @@ class RewardCfg: model_factory: Callable[[Any], torch.nn.Sequential] = mlp_factory + reward_clip: float = None + '''Clamp raw reward output to [-reward_clip, reward_clip]. None disables clamping. + The WGAN-style critic has no inherent output bound, and without one its scale can + drift arbitrarily over training (observed: mean_demo_reward climbing from ~6 to ~16 + across iterations, Value/mean_q_value swinging from -187 to +164) โ€” this bounds how + far that drift can propagate into the TD target / value loss.''' + @dataclass(kw_only=True) class EnsembleValueCfg: @@ -266,6 +273,15 @@ class DynamicsLearnerCfg: enc_lr_scale: float = None '''Scaling factor for the encoder learning rate relative to the dynamics model learning rate.''' + sigreg_coeff: float = None + '''Coefficient for SIGReg auxiliary regularization loss (see mpail2.dynamics.SIGReg). Disabled when None or <= 0.''' + + sigreg_knots: int = None + '''Number of integration knots used by SIGReg.''' + + sigreg_num_proj: int = None + '''Number of random projection vectors used by SIGReg.''' + @dataclass(kw_only=True) class PolicyLearnerCfg: diff --git a/mpail2/configs/defs.py b/mpail2/configs/defs.py index 01e942b..0d5e996 100644 --- a/mpail2/configs/defs.py +++ b/mpail2/configs/defs.py @@ -14,7 +14,7 @@ #### SHARED CONSTANTS #### OPT = "adam" -LR = 3e-4 +LR = 2e-4 HORIZON = 7 OPT_ITERS = 5 GAMMA = 0.99 @@ -54,6 +54,7 @@ class RewardConfig(cfgs.RewardCfg): "use_layer_norm": False, "disable_output_bias": True, }) + reward_clip: float = None @dataclass(kw_only=True) class EnsembleValueConfig(cfgs.EnsembleValueCfg): @@ -83,7 +84,10 @@ class CNNCoderConfig(cfgs.CNNCoderCfg): model_kwargs: dict = field(default_factory=lambda: { **CNN_MODEL_KWARGS, ## OVERRIDES ## - "override_last_layer_activation": True, + # Camera embedding ends in LayerNorm (normalized, zero-centered), not a + # trailing one-sided SiLU which distorts the geometry LayerNorm just + # established (ported from mpail-research fix-main #18). + "override_last_layer_norm": True, }) @dataclass(kw_only=True) @@ -95,9 +99,9 @@ class ProprioCoderConfig(cfgs.MLPCoderCfg): **MODEL_KWARGS, ## OVERRIDES ## "hidden_dims": [256], - "use_layer_norm": False, + # Hidden LN on (inherited from MODEL_KWARGS), no trailing SiLU on the + # latent slice โ€” same rationale as CNNCoderConfig above. "override_last_layer_norm": True, - "override_last_layer_activation": True, }) model_kwargs:dict = field(default_factory=lambda: { @@ -175,10 +179,10 @@ class PlannerConfig(cfgs.PlannerCfg): temperature: float = 2.0 opt_iters: int = OPT_ITERS - reward_cfg: RewardConfig = RewardConfig() - value_cfg: EnsembleValueConfig = EnsembleValueConfig() - sampling_cfg: PolicySamplingConfig = PolicySamplingConfig() - dynamics_cfg: DynamicsConfig = DynamicsConfig() + reward_cfg: RewardConfig = field(default_factory=RewardConfig) + value_cfg: EnsembleValueConfig = field(default_factory=EnsembleValueConfig) + sampling_cfg: PolicySamplingConfig = field(default_factory=PolicySamplingConfig) + dynamics_cfg: DynamicsConfig = field(default_factory=DynamicsConfig) seed: int = 42 u_per_command: int = 1 @@ -210,7 +214,7 @@ class RewardLearnerConfig(cfgs.RewardLearnerCfg): opt_params: dict = field(default_factory=lambda: OPT_PARAMS) - gp_coeff: float = 0.1 + gp_coeff: float = 5.0 gp_target_gradient: float = 1.0 @@ -223,7 +227,7 @@ class PolicyLearnerConfig(cfgs.PolicyLearnerCfg): max_grad_norm: float = 1.0 - target_entropy: float = -3.0 # -ACTION_DIM + target_entropy: float = -2.0 # encourage exploration; -ACTION_DIM=-5 is too conservative alpha_lr: float = LR @@ -236,12 +240,19 @@ class DynamicsLearnerConfig(cfgs.DynamicsLearnerCfg): opt_params: dict = field(default_factory=lambda: OPT_PARAMS) - enc_lr_scale: float = 0.1 + enc_lr_scale: float = 0.08 rho: float = 0.95 recon_coeff: float = 1.0 # Exact coefficient doesn't matter if not using recon loss + # Conservative starting point (ported from mpail-research fix-main #18, which used + # 0.1 on a sim pick-place task) โ€” start smaller so it doesn't swamp the JEP loss + # before it's been tuned here. + sigreg_coeff: float = 0.02 + sigreg_knots: int = 17 + sigreg_num_proj: int = 1024 + @dataclass(kw_only=True) class LearnerConfig(cfgs.MPAIL2LearnerCfg): @@ -259,13 +270,13 @@ class LearnerConfig(cfgs.MPAIL2LearnerCfg): loss_horizon: int = HORIZON # Match num_timesteps in sampling_cfg # Dynamics - ENABLED for training - dynamics_learner_cfg: cfgs.DynamicsLearnerCfg = DynamicsLearnerConfig() + dynamics_learner_cfg: cfgs.DynamicsLearnerCfg = field(default_factory=DynamicsLearnerConfig) # Reward (adversarial training) - reward_learner_cfg: cfgs.RewardLearnerCfg = RewardLearnerConfig() + reward_learner_cfg: cfgs.RewardLearnerCfg = field(default_factory=RewardLearnerConfig) # Value - value_learner_cfg: cfgs.ValueLearnerCfg = ValueLearnerConfig() + value_learner_cfg: cfgs.ValueLearnerCfg = field(default_factory=ValueLearnerConfig) # Policy Learner - policy_learner_cfg: cfgs.PolicyLearnerCfg = PolicyLearnerConfig() + policy_learner_cfg: cfgs.PolicyLearnerCfg = field(default_factory=PolicyLearnerConfig) diff --git a/mpail2/dynamics.py b/mpail2/dynamics.py index 6e6baa3..81ea369 100644 --- a/mpail2/dynamics.py +++ b/mpail2/dynamics.py @@ -9,6 +9,46 @@ from .configs.cfgs import DynamicsCfg +class SIGReg(torch.nn.Module): + """Sketch Isotropic Gaussian Regularizer. + + Projects a batch of latents onto random 1D directions and penalizes deviation + of each projection's empirical characteristic function from a standard + Gaussian's (via a differentiable normality test, quadrature-integrated over + `knots`). Minimizing this pushes variance to spread across many random + directions instead of concentrating in a few (representation/dimensional + collapse) โ€” ported from mpail-research fix-main #18. + """ + + def __init__(self, knots: int = 17, num_proj: int = 1024): + super().__init__() + self.num_proj = num_proj + + t = torch.linspace(0.0, 3.0, knots, dtype=torch.float32) + dt = 3.0 / (knots - 1) + weights = torch.full((knots,), 2.0 * dt, dtype=torch.float32) + weights[[0, -1]] = dt + window = torch.exp(-t.square() / 2.0) + + self.register_buffer("t", t) + self.register_buffer("phi", window) + self.register_buffer("weights", weights * window) + + def forward(self, proj: torch.Tensor) -> torch.Tensor: + """ + Args: + proj: Tensor of shape (T, B, D). + """ + A = torch.randn(proj.size(-1), self.num_proj, device=proj.device, dtype=proj.dtype) + A = A.div_(A.norm(p=2, dim=0)) + + x_t = (proj @ A).unsqueeze(-1) * self.t.to(dtype=proj.dtype) + err = (x_t.cos().mean(-3) - self.phi.to(dtype=proj.dtype)).square() + err = err + x_t.sin().mean(-3).square() + statistic = (err @ self.weights.to(dtype=proj.dtype)) * proj.size(-2) + return statistic.mean() + + class Dynamics(torch.nn.Module): """ Latent dynamics model. Propagates latent states through a feedforward network. diff --git a/mpail2/envs/real/__init__.py b/mpail2/envs/real/__init__.py index 9eb5b16..3972d01 100644 --- a/mpail2/envs/real/__init__.py +++ b/mpail2/envs/real/__init__.py @@ -64,3 +64,38 @@ ) except ImportError: pass + +try: + from .so101 import ( + ACTION_DIM as SO101_ACTION_DIM, + SO101RealWrapper, + SO101RealEnvArgs, + STATE_DIM as SO101_STATE_DIM, + EE_PROPRIO_DIM as SO101_EE_PROPRIO_DIM, + HOME_POSITION_DEG as SO101_HOME_POSITION_DEG, + EE_LOWER_M as SO101_EE_LOWER_M, + EE_UPPER_M as SO101_EE_UPPER_M, + JOINT_LOWER_DEG as SO101_JOINT_LOWER_DEG, + JOINT_UPPER_DEG as SO101_JOINT_UPPER_DEG, + MAX_EPISODE_STEPS as SO101_MAX_EPISODE_STEPS, + make_so101_env, + ) + + __all__.extend( + [ + "SO101RealWrapper", + "SO101RealEnvArgs", + "SO101_STATE_DIM", + "SO101_ACTION_DIM", + "SO101_EE_PROPRIO_DIM", + "SO101_HOME_POSITION_DEG", + "SO101_EE_LOWER_M", + "SO101_EE_UPPER_M", + "SO101_JOINT_LOWER_DEG", + "SO101_JOINT_UPPER_DEG", + "SO101_MAX_EPISODE_STEPS", + "make_so101_env", + ] + ) +except ImportError: + pass diff --git a/mpail2/envs/real/so101/README.md b/mpail2/envs/real/so101/README.md new file mode 100644 index 0000000..2610cdd --- /dev/null +++ b/mpail2/envs/real/so101/README.md @@ -0,0 +1,203 @@ +# SO-101 real-robot environment + +Real SO-101 (SO-ARM101) 6-DOF arm support for MPAIL2 โ€” ported from the +`mpail-lerobot` fork. Structured to mirror `mpail2/envs/real/franka/`: the +gym-env-facing code lives directly in this package, the hardware-owning gRPC +server lives under `network/` (like Franka's `network/server.py` / +`network/client.py`), and the standalone training/data scripts live under +`training/`. + +Unlike Franka (custom hardware driver in `mpail2/envs/real/franka/hardware/`), +SO-101 drives the arm + cameras through LeRobot's own driver stack +(`SO101Follower`, OpenCV/RealSense camera classes) via a patched copy of +LeRobot's `async_inference` protocol โ€” hence `lerobot_patch/`, which has no +Python role inside `mpail2` itself; it's a deploy target for a separate +`lerobot` clone (see below). + +## Layout + +``` +so101/ + __init__.py # exports: SO101RobotEnv, SO101RealWrapper, SO101RealEnvArgs, make_so101_env, ... + so101_env.py # gym.Env client (gRPC) โ€” analogous to Franka's network/client.py's FrankaClient + env_factory.py # make_so101_env() factory + args dataclass wiring + wrappers.py # SO101RealWrapper (MPAIL-shaped obs/action) + robot_limits.py # joint/EE bounds, home pose, dims, camera specs + ik_utils.py # FK/IK (ikpy + soa.urdf), used by so101_env.py + soa.urdf # arm description used by ik_utils.py + transport/ # gRPC stubs (generated from so101_robot.proto / lerobot's services.proto) + network/ + server.py # gRPC server owning the physical arm + cameras + training/ + demo_recording_server.py # gRPC server for recording (obs, next_obs) demo pairs + train_so101_local.py # main MPAIL2 training entry point (in-process env loop) + convert.py # raw_demos*/*.npz -> demo.pt + convert_lerobot.py # LeRobot-recorded dataset -> demo.pt + replay_demo.py # replay/sanity-check a recorded .npz trajectory + check_encoder_collapse.py # offline diagnostic: encoder latent effective-rank check + lerobot_patch/ + async_inference/ # patched lerobot/src/lerobot/async_inference/ โ€” deploy into your lerobot clone +``` + +## Install + +```bash +conda create -n mpail2 python=3.10 +conda activate mpail2 +pip install -e ".[so101]" +``` + +You'll also need a separate `lerobot` conda env (Python 3.12) with the +[`lerobot`](https://github.com/huggingface/lerobot) repo installed โ€” it owns +the low-level Feetech servo + camera drivers this package talks to: + +```bash +conda create -n lerobot python=3.12 +conda activate lerobot +cd && pip install -e . +pip install pyrealsense2 # RealSense camera support +pip install -e "[so101]" # so mpail2.envs.real.so101.network.server is importable here too +``` + +Deploy the LeRobot patch once (re-run after pulling upstream `lerobot` +changes, or after editing anything under `lerobot_patch/async_inference/` +here โ€” it fixes a camera-reconnect bug, a timestep off-by-one, and adds a +block-until-server-ready sync point between episodes/training updates): + +```bash +cp -r mpail2/envs/real/so101/lerobot_patch/async_inference/* /src/lerobot/async_inference/ +``` + +### Known issue: the `lerobot` env needs mpail2's *full* dependency stack + +Because `network/server.py` is a submodule of the `mpail2` package, importing +it (even via `-m`) runs `mpail2/__init__.py` first โ€” which eagerly imports the +whole algorithm stack (`configs`, `encoder`, `learner`, ...), pulling in +`torch`/`hydra-core`/`wandb`/`scipy`/`matplotlib`, not just the +grpc/opencv/pyrealsense2 that `network/server.py` itself actually needs. This +is inherent to nesting the server inside the package (Franka's +`network/server.py` has the exact same coupling) โ€” it isn't specific to +SO-101. Concretely, this means the `lerobot` conda env needs a `matplotlib` +version compatible with `mpail2/utils/rollout_vis.py`'s `matplotlib.cm.get_cmap` +usage (removed in matplotlib โ‰ฅ 3.9) โ€” pin an older matplotlib there if you hit +`ImportError: cannot import name 'get_cmap' from 'matplotlib.cm'`. + +## Running scripts + +Everything here is a proper submodule of the installed `mpail2` package, so +invoke with `-m` from anywhere (no need to `cd` to a particular directory, +unlike the original `mpail-lerobot` fork's top-level scripts): + +```bash +python -m mpail2.envs.real.so101.network.server --help +python -m mpail2.envs.real.so101.training.train_so101_local --help +``` + +## Workflow + +### 1. Calibrate (once per arm, `lerobot` env) + +```bash +conda activate lerobot +lerobot-find-port +lerobot-calibrate --robot.type=so_follower --robot.port=/dev/ttyACM0 --robot.id= +lerobot-calibrate --teleop.type=so_leader --teleop.port=/dev/ttyACM1 --teleop.id= +``` + +### 2. Start the robot-side server (`lerobot` env, keep running) + +```bash +conda activate lerobot +python -m mpail2.envs.real.so101.network.server \ + --robot_port /dev/ttyACM0 --robot_id \ + --cam_index /dev/video0 --cam2_serial --grpc_port 7070 +``` + +Owns the arm + both cameras; backs `training/train_so101_local.py`'s online +training loop. See its `--help` for servo-tuning flags +(`--p_coefficient`, `--i_coefficient`, `--goal_velocity`, ...) if motion is +shaky or not settling. + +### 3. Collect demonstrations + +Start the recording server (`mpail2` env): + +```bash +conda activate mpail2 +python -m mpail2.envs.real.so101.training.demo_recording_server --collect_dir ./raw_demos2 --flush_every 200 +``` + +It holds the arm still (or drives home between episodes) and records every +`(obs_t, obs_t+1)` pair it receives. Drive the arm via LeRobot's standard +teleop client (`lerobot` env, separate terminal): + +```bash +python -m lerobot.async_inference.robot_client \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyACM0 \ + --robot.id= \ + --robot.cameras="{cam: {type: opencv, index_or_path: /dev/video0, width: 640, height: 480, fps: 30}, cam2: {type: intelrealsense, serial_number_or_name: , width: 640, height: 480, fps: 30}}" \ + --teleop.type=so100_leader \ + --teleop.port=/dev/ttyACM1 \ + --teleop.id= \ + --server_address=127.0.0.1:8080 \ + --policy_type=act \ + --pretrained_name_or_path=dummy \ + --actions_per_chunk=1 \ + --task="pick up the cup" +``` + +Flag notes: +- `--robot.cameras`: `cam` (wrist, OpenCV) and `cam2` (RealSense, by serial number). +- `--server_address`: must match `demo_recording_server.py`'s `--port` (default 8080). +- `--policy_type` / `--pretrained_name_or_path`: required by `robot_client`'s + CLI even here, where the server ignores them and just echoes joint state + back โ€” `act` / `dummy` are placeholders, not a real policy. +- Move the leader arm to demonstrate the task; the follower mirrors it and + every step gets recorded. Each episode auto-ends (saves, homes, pauses, + resumes) after `--max_episode_steps` (default 200) steps. + +### 4. Convert to training format + +```bash +conda activate mpail2 +python -m mpail2.envs.real.so101.training.convert --dirs raw_demos2 --out demo.pt --img_w 64 --img_h 48 +``` + +(Recorded via a LeRobot dataset instead? Use `convert_lerobot.py` in the +`lerobot` env instead.) + +Sanity-check a trajectory by replaying it on the real arm: + +```bash +conda activate lerobot +python -m mpail2.envs.real.so101.training.replay_demo raw_demos2/traj_0000.npz --port /dev/ttyACM0 --robot_id +``` + +### 5. Train + +With the robot-side server (step 2) still running: + +```bash +conda activate mpail2 +python -m mpail2.envs.real.so101.training.train_so101_local \ + --demo_path demo.pt --robot_host 127.0.0.1 --robot_port 7070 \ + --device cuda --speed_scale 0.4 --lpf_alpha 0.5 --wandb +``` + +See `--help` for the full flag list (MPPI sampling, gripper hold steps, +checkpointing, eval mode, ...). `--eval --load_checkpoint ` rolls out a +trained checkpoint with no training updates; add `--eval_policy_only` to +bypass CEM/MPPI and use the policy network's own deterministic action. + +### 6. Diagnose encoder collapse (optional) + +```bash +conda activate mpail2 +python -m mpail2.envs.real.so101.training.check_encoder_collapse \ + --demo_path demo.pt --checkpoint logs/so101_local/models/model_N.pt +``` + +Reports per-dimension latent std and effective rank (participation ratio) โ€” a +low effective rank relative to `latent_dim` is the signature of +representation collapse even when the JEP/dynamics loss looks fine. diff --git a/mpail2/envs/real/so101/__init__.py b/mpail2/envs/real/so101/__init__.py new file mode 100644 index 0000000..91a39f6 --- /dev/null +++ b/mpail2/envs/real/so101/__init__.py @@ -0,0 +1,57 @@ +"""SO-101 (SO-ARM101) real-robot environment for MPAIL.""" + +from .robot_limits import ( + ACTION_DIM, + CONTROL_FREQUENCY_HZ, + DEFAULT_ROBOT_HOST, + DEFAULT_ROBOT_PORT, + HOME_EE_M, + HOME_GRIPPER_DEG, + HOME_POSITION_DEG, + EE_HALF_RANGE_M, + EE_LOWER_M, + EE_UPPER_M, + GRIPPER_HALF_RANGE, + GRIPPER_LOWER_DEG, + GRIPPER_UPPER_DEG, + JOINT_LOWER_DEG, + JOINT_NAMES, + JOINT_UPPER_DEG, + MAX_DELTA_DEG, + MAX_DELTA_M, + MAX_EPISODE_STEPS, + STATE_DIM, + EE_PROPRIO_DIM, +) +from .so101_env import OBS_KEY, SO101RobotEnv +from .wrappers import SO101RealWrapper +from .env_factory import SO101RealEnvArgs, make_so101_env + +__all__ = [ + "SO101RobotEnv", + "SO101RealWrapper", + "SO101RealEnvArgs", + "make_so101_env", + "OBS_KEY", + "STATE_DIM", + "EE_PROPRIO_DIM", + "ACTION_DIM", + "MAX_DELTA_DEG", + "MAX_DELTA_M", + "MAX_EPISODE_STEPS", + "CONTROL_FREQUENCY_HZ", + "JOINT_NAMES", + "JOINT_LOWER_DEG", + "JOINT_UPPER_DEG", + "HOME_POSITION_DEG", + "HOME_EE_M", + "HOME_GRIPPER_DEG", + "EE_HALF_RANGE_M", + "EE_LOWER_M", + "EE_UPPER_M", + "GRIPPER_HALF_RANGE", + "GRIPPER_LOWER_DEG", + "GRIPPER_UPPER_DEG", + "DEFAULT_ROBOT_HOST", + "DEFAULT_ROBOT_PORT", +] diff --git a/mpail2/envs/real/so101/env_factory.py b/mpail2/envs/real/so101/env_factory.py new file mode 100644 index 0000000..6e17094 --- /dev/null +++ b/mpail2/envs/real/so101/env_factory.py @@ -0,0 +1,50 @@ +"""Factory for the SO-101 real-robot stack (mirrors kinova.env_factory).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .robot_limits import ( + ACTION_DIM, + CONTROL_FREQUENCY_HZ, + DEFAULT_ROBOT_HOST, + DEFAULT_ROBOT_PORT, + MAX_EPISODE_STEPS, + STATE_DIM, +) +from .so101_env import SO101RobotEnv +from .wrappers import SO101RealWrapper + + +@dataclass +class SO101RealEnvArgs: + device: str = "cuda" + host: str = DEFAULT_ROBOT_HOST + port: int = DEFAULT_ROBOT_PORT + control_hz: float = CONTROL_FREQUENCY_HZ + max_episode_length: int = MAX_EPISODE_STEPS + state_dim: int = STATE_DIM + action_dim: int = ACTION_DIM + mock: bool = False + speed_scale: float = 1.0 + lpf_alpha: float = 1.0 + reset_pause_seconds: float = 3.0 + gripper_hold_steps: int = 1 + + +def make_so101_env(args: SO101RealEnvArgs | None = None) -> SO101RealWrapper: + """Build a fully wrapped SO-101 env ready for MPAIL2Runner.""" + if args is None: + args = SO101RealEnvArgs() + base = SO101RobotEnv( + host=args.host, + port=args.port, + control_frequency=args.control_hz, + max_episode_steps=args.max_episode_length, + mock=args.mock, + speed_scale=args.speed_scale, + lpf_alpha=args.lpf_alpha, + reset_pause_seconds=args.reset_pause_seconds, + gripper_hold_steps=args.gripper_hold_steps, + ) + return SO101RealWrapper(base, device=args.device) diff --git a/mpail2/envs/real/so101/ik_utils.py b/mpail2/envs/real/so101/ik_utils.py new file mode 100644 index 0000000..d87fd7f --- /dev/null +++ b/mpail2/envs/real/so101/ik_utils.py @@ -0,0 +1,167 @@ +"""FK / IK utilities for the SO-101 arm using the SOA URDF and ikpy. + +Chain layout (7 links): + 0 Base link โ€” fixed, inactive + 1 shoulder_pan + 2 shoulder_lift + 3 elbow_flex + 4 wrist_flex + 5 wrist_roll + 6 gripper โ€” not part of FK chain end-effector, inactive + +IK method: Jacobian pseudo-inverse (differential IK), same principle as Franka's +internal Cartesian controller. Iterates: ฮ”q = Jโบ ฮ”x until convergence or max_iter. +Uses damped least-squares (DLS) near singularities for numerical stability. + +wrist_roll is excluded from the Jacobian โ€” it does not affect EE position. Both +so101_env.py and grpc_policy_server.py use a 5-dim [x, y, z, wrist_roll, gripper] +action and set wrist_roll directly from index 3 (delta-style); so101_env.py's +xyz is an absolute position target, grpc_policy_server.py's is still a delta โ€” +only that difference remains between the two. +""" + +import os + +import numpy as np + +# Resolved relative to this file (not hardcoded to any one clone's absolute path) +# so it keeps working regardless of where this repo is checked out. +URDF_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "soa.urdf") + +# Full chain mask โ€” all joints active for FK (wrist_roll included so FK is correct). +_ACTIVE_MASK_FK = [False, True, True, True, True, True, False] + +# Position-controlling joints only (indices into the 5-elem arm vector). +# wrist_roll (index 4) is excluded โ€” it does not move the EE tip. +_POS_JOINTS = [0, 1, 2, 3] # shoulder_pan, shoulder_lift, elbow_flex, wrist_flex + +_chain = None + + +def _get_chain(): + global _chain + if _chain is None: + import ikpy.chain + import warnings + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _chain = ikpy.chain.Chain.from_urdf_file(URDF_PATH, active_links_mask=_ACTIVE_MASK_FK) + return _chain + + +def fk(arm_joints_deg: np.ndarray) -> np.ndarray: + """Forward kinematics: 5 arm joint angles (degrees) โ†’ EE xyz (metres). + + Args: + arm_joints_deg: shape (5,) โ€” shoulder_pan โ€ฆ wrist_roll in degrees. + + Returns: + xyz: shape (3,) float32 in metres. + """ + chain = _get_chain() + rad = np.deg2rad(arm_joints_deg.astype(np.float64)) + full = [0.0] + list(rad) + [0.0] + T = chain.forward_kinematics(full) + return T[:3, 3].astype(np.float32) + + +def fk_pose(arm_joints_deg: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Forward kinematics: 5 arm joint angles (degrees) โ†’ EE xyz + quaternion. + + Args: + arm_joints_deg: shape (5,) โ€” shoulder_pan โ€ฆ wrist_roll in degrees. + + Returns: + xyz: shape (3,) float32 in metres. + quat: shape (4,) float32 โ€” [x, y, z, w] scalar-last convention. + """ + from scipy.spatial.transform import Rotation as R + chain = _get_chain() + rad = np.deg2rad(arm_joints_deg.astype(np.float64)) + full = [0.0] + list(rad) + [0.0] + T = chain.forward_kinematics(full) + xyz = T[:3, 3].astype(np.float32) + quat = R.from_matrix(T[:3, :3]).as_quat(scalar_first=False).astype(np.float32) # [x,y,z,w] + return xyz, quat + + +def joints_to_ee_proprio(joint_state: np.ndarray) -> np.ndarray: + """Convert 6-dim joint state (degrees) to 13-dim proprioception. + + Returns [ee_x, ee_y, ee_z, qx, qy, qz, qw, j0..j5] โ€” matches Franka convention. + Shared by grpc_policy_server.py and mpail2/envs/real/so101/so101_env.py so both + build the exact same observation the trained encoder expects. + """ + xyz, quat = fk_pose(joint_state[:5]) + return np.concatenate([xyz, quat, joint_state], dtype=np.float32) + + +def _jacobian(arm_joints_deg: np.ndarray, eps_deg: float = 0.5) -> np.ndarray: + """Numerical position Jacobian d(xyz)/d(q) for the 4 position-controlling joints. + + Returns shape (3, 4): columns are shoulder_pan, shoulder_lift, elbow_flex, wrist_flex. + Units: metres per radian. + """ + xyz0 = fk(arm_joints_deg).astype(np.float64) + eps_rad = np.deg2rad(eps_deg) + J = np.zeros((3, 4), dtype=np.float64) + for col, joint_idx in enumerate(_POS_JOINTS): + q_p = arm_joints_deg.copy().astype(np.float64) + q_p[joint_idx] += eps_deg + J[:, col] = (fk(q_p).astype(np.float64) - xyz0) / eps_rad + return J + + +def ik( + xyz_target: np.ndarray, + initial_arm_deg: np.ndarray | None = None, + max_iter: int = 8, + tol_m: float = 0.002, + damping: float = 0.05, +) -> np.ndarray: + """Differential IK: Jacobian pseudo-inverse, same principle as Franka's controller. + + Iteratively steps joints along the Jacobian direction until the EE reaches + xyz_target within tol_m, or max_iter is exhausted. Uses damped least-squares + (DLS) for stability near singularities. + + wrist_roll (index 4) is never modified here โ€” both callers (so101_env.py, + grpc_policy_server.py) set it themselves from their own action's wrist_roll + component after calling this function. + + Args: + xyz_target: shape (3,) desired EE position in metres. + initial_arm_deg: shape (5,) current joint angles in degrees (warm-start). + max_iter: maximum Jacobian iterations (default 8). + tol_m: convergence threshold in metres (default 2 mm). + damping: DLS damping factor ฮป โ€” higher = more stable, less accurate. + + Returns: + arm_joints_deg: shape (5,) float32 โ€” shoulder_pan โ€ฆ wrist_roll. + """ + if initial_arm_deg is None: + initial_arm_deg = np.array([-21.19, -5.41, 1.58, 99.47, -14.20], np.float32) + + q = initial_arm_deg.copy().astype(np.float64) + target = xyz_target.astype(np.float64) + + for _ in range(max_iter): + error = target - fk(q).astype(np.float64) + if np.linalg.norm(error) < tol_m: + break + + J = _jacobian(q) # (3, 4) + # Damped least-squares: J^T (J J^T + ฮปยฒI)^{-1} + JJT = J @ J.T # (3, 3) + J_dls = J.T @ np.linalg.inv(JJT + damping ** 2 * np.eye(3)) # (4, 3) + + dq_rad = J_dls @ error # (4,) in radians + dq_deg = np.rad2deg(dq_rad) + + # Clamp per-step joint movement to avoid large jumps + dq_deg = np.clip(dq_deg, -8.0, 8.0) + + for col, joint_idx in enumerate(_POS_JOINTS): + q[joint_idx] += dq_deg[col] + + return q.astype(np.float32) diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/__init__.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/__init__.py new file mode 100644 index 0000000..8d7a225 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Async inference server/client. + +Requires: ``pip install 'lerobot[async]'`` + +Available modules (import directly):: + + from lerobot.async_inference.policy_server import ... + from lerobot.async_inference.robot_client import ... +""" + +from lerobot.utils.import_utils import require_package + +require_package("grpcio", extra="async", import_name="grpc") + +__all__: list[str] = [] diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/configs.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/configs.py new file mode 100644 index 0000000..67c43f5 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/configs.py @@ -0,0 +1,252 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Optional + +import torch + +from lerobot.robots.config import RobotConfig +from lerobot.teleoperators.config import TeleoperatorConfig + +from .constants import ( + DEFAULT_FPS, + DEFAULT_INFERENCE_LATENCY, + DEFAULT_OBS_QUEUE_TIMEOUT, +) + +# Aggregate function registry for CLI usage +AGGREGATE_FUNCTIONS = { + "weighted_average": lambda old, new: 0.3 * old + 0.7 * new, + "latest_only": lambda old, new: new, + "average": lambda old, new: 0.5 * old + 0.5 * new, + "conservative": lambda old, new: 0.7 * old + 0.3 * new, +} + + +def get_aggregate_function(name: str) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]: + """Get aggregate function by name from registry.""" + if name not in AGGREGATE_FUNCTIONS: + available = list(AGGREGATE_FUNCTIONS.keys()) + raise ValueError(f"Unknown aggregate function '{name}'. Available: {available}") + return AGGREGATE_FUNCTIONS[name] + + +@dataclass +class PolicyServerConfig: + """Configuration for PolicyServer. + + This class defines all configurable parameters for the PolicyServer, + including networking settings and action chunking specifications. + """ + + # Networking configuration + host: str = field(default="localhost", metadata={"help": "Host address to bind the server to"}) + port: int = field(default=8080, metadata={"help": "Port number to bind the server to"}) + + # Timing configuration + fps: int = field(default=DEFAULT_FPS, metadata={"help": "Frames per second"}) + inference_latency: float = field( + default=DEFAULT_INFERENCE_LATENCY, metadata={"help": "Target inference latency in seconds"} + ) + + obs_queue_timeout: float = field( + default=DEFAULT_OBS_QUEUE_TIMEOUT, metadata={"help": "Timeout for observation queue in seconds"} + ) + + def __post_init__(self): + """Validate configuration after initialization.""" + if self.port < 1 or self.port > 65535: + raise ValueError(f"Port must be between 1 and 65535, got {self.port}") + + if self.environment_dt <= 0: + raise ValueError(f"environment_dt must be positive, got {self.environment_dt}") + + if self.inference_latency < 0: + raise ValueError(f"inference_latency must be non-negative, got {self.inference_latency}") + + if self.obs_queue_timeout < 0: + raise ValueError(f"obs_queue_timeout must be non-negative, got {self.obs_queue_timeout}") + + @classmethod + def from_dict(cls, config_dict: dict) -> "PolicyServerConfig": + """Create a PolicyServerConfig from a dictionary.""" + return cls(**config_dict) + + @property + def environment_dt(self) -> float: + """Environment time step, in seconds""" + return 1 / self.fps + + def to_dict(self) -> dict: + """Convert the configuration to a dictionary.""" + return { + "host": self.host, + "port": self.port, + "fps": self.fps, + "environment_dt": self.environment_dt, + "inference_latency": self.inference_latency, + } + + +@dataclass +class RobotClientConfig: + """Configuration for RobotClient. + + This class defines all configurable parameters for the RobotClient, + including network connection, policy settings, and control behavior. + """ + + # Policy configuration + policy_type: str = field(metadata={"help": "Type of policy to use"}) + pretrained_name_or_path: str = field(metadata={"help": "Pretrained model name or path"}) + + # Robot configuration (for CLI usage - robot instance will be created from this) + robot: RobotConfig = field(metadata={"help": "Robot configuration"}) + + # Policies typically output K actions at max, but we can use less to avoid wasting bandwidth (as actions + # would be aggregated on the client side anyway, depending on the value of `chunk_size_threshold`) + actions_per_chunk: int = field(metadata={"help": "Number of actions per chunk"}) + + # Task instruction for the robot to execute (e.g., 'fold my tshirt') + task: str = field(default="", metadata={"help": "Task instruction for the robot to execute"}) + + # Network configuration + server_address: str = field(default="localhost:8080", metadata={"help": "Server address to connect to"}) + + # Device configuration + policy_device: str = field(default="cpu", metadata={"help": "Device for policy inference"}) + client_device: str = field( + default="cpu", + metadata={ + "help": "Device to move actions to after receiving from server (e.g., for downstream planners)" + }, + ) + + # Control behavior configuration + chunk_size_threshold: float = field(default=0.5, metadata={"help": "Threshold for chunk size control"}) + fps: int = field(default=DEFAULT_FPS, metadata={"help": "Frames per second"}) + + # Aggregate function configuration (CLI-compatible) + aggregate_fn_name: str = field( + default="weighted_average", + metadata={"help": f"Name of aggregate function to use. Options: {list(AGGREGATE_FUNCTIONS.keys())}"}, + ) + + # Debug configuration + debug_visualize_queue_size: bool = field( + default=False, metadata={"help": "Visualize the action queue size"} + ) + + # โ”€โ”€ Teleoperation (leader arm) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # When set, actions come from the leader arm instead of the policy server. + # Observations (including the teleop action) are still streamed to the server + # so planner_server.py can record them. policy_type / pretrained_name_or_path + # are still required for the gRPC handshake but are ignored server-side. + teleop: Optional[TeleoperatorConfig] = field( + default=None, + metadata={"help": "Leader arm config. If set, enables teleoperation mode."}, + ) + + # โ”€โ”€ Home position / smooth reset (mirrors HTTP robot_server capabilities) โ”€โ”€ + # Space-separated joint positions in degrees, matching robot.action_features order. + # Empty string = no reset (original behaviour). + home_joints: str = field( + default="", + metadata={"help": "Space-separated home joint positions in degrees. Empty = skip reset."}, + ) + # Number of linear interpolation steps when moving to home. + reset_steps: int = field( + default=80, + metadata={"help": "Interpolation steps when moving to home position (more = slower)."}, + ) + # Frequency of each interpolation step. + reset_hz: float = field( + default=50.0, + metadata={"help": "Hz of each interpolation step during reset."}, + ) + + # โ”€โ”€ Episode mode โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # If set, the control loop stops after this many executed actions and resets. + # None = run indefinitely (original behaviour). + max_episode_steps: int | None = field( + default=None, + metadata={"help": "Actions per episode. None = run forever (original behaviour)."}, + ) + # Total number of episodes to run before exiting. None = run forever. + num_episodes: int | None = field( + default=None, + metadata={"help": "Number of episodes. None = run forever (original behaviour)."}, + ) + # Pause between episodes (after server training, before moving to home). + # During this window the robot holds still and no images are sent. + reset_pause_seconds: float = field( + default=0.0, + metadata={"help": "Seconds to pause between episodes for scene reset. Default 0 (no pause)."}, + ) + + @property + def environment_dt(self) -> float: + """Environment time step, in seconds""" + return 1 / self.fps + + def __post_init__(self): + """Validate configuration after initialization.""" + if not self.server_address: + raise ValueError("server_address cannot be empty") + + if not self.policy_type: + raise ValueError("policy_type cannot be empty") + + if not self.pretrained_name_or_path: + raise ValueError("pretrained_name_or_path cannot be empty") + + if not self.policy_device: + raise ValueError("policy_device cannot be empty") + + if not self.client_device: + raise ValueError("client_device cannot be empty") + + if self.chunk_size_threshold < 0 or self.chunk_size_threshold > 1: + raise ValueError(f"chunk_size_threshold must be between 0 and 1, got {self.chunk_size_threshold}") + + if self.fps <= 0: + raise ValueError(f"fps must be positive, got {self.fps}") + + if self.actions_per_chunk <= 0: + raise ValueError(f"actions_per_chunk must be positive, got {self.actions_per_chunk}") + + self.aggregate_fn = get_aggregate_function(self.aggregate_fn_name) + + @classmethod + def from_dict(cls, config_dict: dict) -> "RobotClientConfig": + """Create a RobotClientConfig from a dictionary.""" + return cls(**config_dict) + + def to_dict(self) -> dict: + """Convert the configuration to a dictionary.""" + return { + "server_address": self.server_address, + "policy_type": self.policy_type, + "pretrained_name_or_path": self.pretrained_name_or_path, + "policy_device": self.policy_device, + "client_device": self.client_device, + "chunk_size_threshold": self.chunk_size_threshold, + "fps": self.fps, + "actions_per_chunk": self.actions_per_chunk, + "task": self.task, + "debug_visualize_queue_size": self.debug_visualize_queue_size, + "aggregate_fn_name": self.aggregate_fn_name, + } diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/constants.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/constants.py new file mode 100644 index 0000000..56910e6 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/constants.py @@ -0,0 +1,29 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Client side: The environment evolves with a time resolution equal to 1/fps""" + +DEFAULT_FPS = 30 + +"""Server side: Running inference on (at most) 1/fps""" +DEFAULT_INFERENCE_LATENCY = 1 / DEFAULT_FPS + +"""Server side: Timeout for observation queue in seconds""" +DEFAULT_OBS_QUEUE_TIMEOUT = 2 + +# All action chunking policies +SUPPORTED_POLICIES = ["act", "smolvla", "diffusion", "tdmpc", "vqbet", "pi0", "pi05", "groot"] + +# TODO: Add all other robots +SUPPORTED_ROBOTS = ["so100_follower", "so101_follower", "bi_so_follower", "omx_follower"] diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/external_planner_server.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/external_planner_server.py new file mode 100644 index 0000000..18c536d --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/external_planner_server.py @@ -0,0 +1,359 @@ +# ExternalPlannerServer +# +# A drop-in replacement for PolicyServer that, instead of running a local neural-network policy, +# forwards every observation to an EXTERNAL planner over HTTP (JSON) and returns whatever +# actions that planner sends back. +# +# Usage: +# python -m lerobot.async_inference.external_planner_server \ +# --host=0.0.0.0 \ +# --port=8080 \ +# --planner_url=http://192.168.1.50:9090/plan +# +# The robot-side RobotClient is UNCHANGED โ€” point it at this server's host:port as usual. + +import dataclasses +import json +import logging +import pickle # nosec +import threading +import time +from concurrent import futures +from queue import Empty, Queue + +import grpc +import numpy as np +import requests + +from lerobot.transport import services_pb2, services_pb2_grpc +from lerobot.transport.utils import receive_bytes_in_chunks + +from .helpers import TimedAction, TimedObservation, get_logger + +logger = get_logger("external_planner_server") + + +# --------------------------------------------------------------------------- +# What is SENT to the external planner (HTTP POST body, JSON) +# --------------------------------------------------------------------------- +# +# { +# "timestep": int, -- which control step this observation belongs to +# "timestamp": float, -- unix time (seconds) when the robot captured this obs +# "must_go": bool, -- True when the robot's action queue is empty and needs actions NOW +# "observation": { +# "task": "string", -- task description passed from robot client +# "shoulder_pan.pos": float, -- raw motor position keys, one per joint +# "shoulder_lift.pos": float, -- key names match robot.action_features exactly +# "elbow_flex.pos": float, -- (actual names depend on your robot model) +# "wrist_flex.pos": float, +# "wrist_roll.pos": float, +# "gripper.pos": float, +# "front": [...], -- camera image as nested list (H, W, 3) uint8, +# -- only included when --include_images is set +# ... (one key per camera) +# } +# } +# +# IMPORTANT: the order of motor keys in "observation" matches robot.action_features. +# Your planner must return actions in that SAME order. +# +# --------------------------------------------------------------------------- +# What is EXPECTED back from the external planner (HTTP response body, JSON) +# --------------------------------------------------------------------------- +# +# { +# "actions": [ +# [float, float, ..., float], -- action vector for timestep N (one float per joint, +# [float, float, ..., float], -- in the same order as the motor keys above) +# ... -- return as many steps as your planner wants +# ] +# } +# +# Each inner list must have exactly action_dim floats (= number of motors). +# The first element is executed at the CURRENT timestep; the rest are queued in order. +# --------------------------------------------------------------------------- + + +def _timed_obs_to_json(obs: TimedObservation, include_images: bool = False) -> dict: + """Serialize a TimedObservation to a plain JSON-safe dict.""" + raw = obs.get_observation() + observation_payload = {} + + for key, value in raw.items(): + if key == "task": + observation_payload["task"] = value + continue + + if isinstance(value, np.ndarray): + if not include_images and value.ndim == 3: + # Skip image tensors by default to keep payload small. + # Set include_images=True if your planner needs them. + continue + observation_payload[key] = value.tolist() + elif hasattr(value, "tolist"): + observation_payload[key] = value.tolist() + else: + observation_payload[key] = value + + return { + "timestep": obs.get_timestep(), + "timestamp": obs.get_timestamp(), + "must_go": bool(obs.must_go), + "observation": observation_payload, + } + + +def _json_to_timed_actions( + response: dict, + base_timestep: int, + base_timestamp: float, + environment_dt: float, +) -> list[TimedAction]: + """Parse the planner's JSON response into a list of TimedAction.""" + import torch + + raw_actions = response["actions"] # list[list[float]] + timed_actions = [] + for i, action_vec in enumerate(raw_actions): + timed_actions.append( + TimedAction( + timestamp=base_timestamp + i * environment_dt, + timestep=base_timestep + i, + action=torch.tensor(action_vec, dtype=torch.float32), + ) + ) + return timed_actions + + +class ExternalPlannerServer(services_pb2_grpc.AsyncInferenceServicer): + """gRPC AsyncInference servicer that delegates planning to an external HTTP server.""" + + def __init__( + self, + planner_url: str, + environment_dt: float = 1.0 / 30, + obs_queue_timeout: float = 1.0, + include_images: bool = False, + http_timeout: float = 5.0, + ): + """ + Args: + planner_url: Full URL of the external planner endpoint, + e.g. "http://192.168.1.50:9090/plan" + environment_dt: Control period in seconds (used to timestamp returned actions). + obs_queue_timeout: How long GetActions blocks waiting for an observation (seconds). + include_images: Whether to include camera images in the JSON payload. + Images are large; only enable if your planner needs them. + http_timeout: Requests timeout for the planner HTTP call (seconds). + """ + self.planner_url = planner_url + self.environment_dt = environment_dt + self.obs_queue_timeout = obs_queue_timeout + self.include_images = include_images + self.http_timeout = http_timeout + + self.shutdown_event = threading.Event() + self.observation_queue: Queue[TimedObservation] = Queue(maxsize=1) + + @property + def running(self): + return not self.shutdown_event.is_set() + + # ------------------------------------------------------------------ + # gRPC handlers (called by the robot's RobotClient) + # ------------------------------------------------------------------ + + def Ready(self, request, context): # noqa: N802 + """Handshake: robot checks the server is alive.""" + logger.info(f"Robot client connected: {context.peer()}") + self.shutdown_event.clear() + self.observation_queue = Queue(maxsize=1) + return services_pb2.Empty() + + def SendPolicyInstructions(self, request, context): # noqa: N802 + """Robot sends policy metadata on startup. + + We log it but do not load any local model โ€” the external planner owns inference. + If your external planner needs these details, forward them here. + """ + policy_specs = pickle.loads(request.data) # nosec + logger.info( + f"Policy instructions received (not used locally): " + f"policy_type={policy_specs.policy_type}, " + f"pretrained={policy_specs.pretrained_name_or_path}, " + f"actions_per_chunk={policy_specs.actions_per_chunk}" + ) + return services_pb2.Empty() + + def SendObservations(self, request_iterator, context): # noqa: N802 + """Robot streams its latest observation here. + + The observation is deserialized and placed in the single-slot queue, + replacing any stale observation that has not been processed yet. + """ + received_bytes = receive_bytes_in_chunks( + request_iterator, None, self.shutdown_event, logger + ) + timed_obs: TimedObservation = pickle.loads(received_bytes) # nosec + + logger.debug(f"Received observation #{timed_obs.get_timestep()} (must_go={timed_obs.must_go})") + + # Replace stale observation if queue is full + if self.observation_queue.full(): + try: + self.observation_queue.get_nowait() + except Empty: + pass + + self.observation_queue.put(timed_obs) + return services_pb2.Empty() + + def GetActions(self, request, context): # noqa: N802 + """Robot polls here to receive its next action chunk. + + 1. Wait for the latest observation. + 2. Forward it to the external planner as JSON. + 3. Parse the planner's response and return it to the robot. + """ + try: + obs: TimedObservation = self.observation_queue.get(timeout=self.obs_queue_timeout) + except Empty: + logger.debug("No observation available yet, returning empty actions.") + return services_pb2.Actions(data=b"") + + logger.info(f"Forwarding observation #{obs.get_timestep()} to external planner at {self.planner_url}") + + payload = _timed_obs_to_json(obs, include_images=self.include_images) + + try: + t0 = time.perf_counter() + response = requests.post( + self.planner_url, + data=json.dumps(payload), + headers={"Content-Type": "application/json"}, + timeout=self.http_timeout, + ) + response.raise_for_status() + elapsed = time.perf_counter() - t0 + logger.info(f"External planner responded in {elapsed * 1000:.1f}ms") + except requests.RequestException as e: + logger.error(f"External planner request failed: {e}") + return services_pb2.Actions(data=b"") + + timed_actions = _json_to_timed_actions( + response.json(), + base_timestep=obs.get_timestep(), + base_timestamp=obs.get_timestamp(), + environment_dt=self.environment_dt, + ) + + actions_bytes = pickle.dumps(timed_actions) # nosec + return services_pb2.Actions(data=actions_bytes) + + def Plan(self, request_iterator, context): # noqa: N802 + """Direct synchronous plan RPC: receive observation stream, call planner, return actions. + + Unlike the SendObservations + GetActions two-step, this is a single blocking call: + the client streams observation chunks, the server forwards them to the HTTP planner, + and returns the resulting actions directly in the response. + """ + received_bytes = receive_bytes_in_chunks( + request_iterator, None, self.shutdown_event, logger + ) + if not received_bytes: + return services_pb2.Actions(data=b"") + + timed_obs: TimedObservation = pickle.loads(received_bytes) # nosec + logger.debug(f"Plan RPC: observation #{timed_obs.get_timestep()} (must_go={timed_obs.must_go})") + + payload = _timed_obs_to_json(timed_obs, include_images=self.include_images) + + try: + t0 = time.perf_counter() + response = requests.post( + self.planner_url, + data=json.dumps(payload), + headers={"Content-Type": "application/json"}, + timeout=self.http_timeout, + ) + response.raise_for_status() + elapsed = time.perf_counter() - t0 + logger.info(f"Plan RPC: planner responded in {elapsed * 1000:.1f}ms for obs #{timed_obs.get_timestep()}") + except requests.RequestException as e: + logger.error(f"Plan RPC: planner request failed: {e}") + return services_pb2.Actions(data=b"") + + timed_actions = _json_to_timed_actions( + response.json(), + base_timestep=timed_obs.get_timestep(), + base_timestamp=timed_obs.get_timestamp(), + environment_dt=self.environment_dt, + ) + + actions_bytes = pickle.dumps(timed_actions) # nosec + return services_pb2.Actions(data=actions_bytes) + + def stop(self): + self.shutdown_event.set() + logger.info("ExternalPlannerServer stopped.") + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +def serve( + host: str = "0.0.0.0", + port: int = 8080, + planner_url: str = "http://localhost:9090/plan", + environment_dt: float = 1.0 / 30, + obs_queue_timeout: float = 1.0, + include_images: bool = False, + http_timeout: float = 5.0, +): + server_instance = ExternalPlannerServer( + planner_url=planner_url, + environment_dt=environment_dt, + obs_queue_timeout=obs_queue_timeout, + include_images=include_images, + http_timeout=http_timeout, + ) + + grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + services_pb2_grpc.add_AsyncInferenceServicer_to_server(server_instance, grpc_server) + grpc_server.add_insecure_port(f"{host}:{port}") + + logger.info(f"ExternalPlannerServer listening on {host}:{port}") + logger.info(f"Forwarding observations to external planner at: {planner_url}") + grpc_server.start() + + try: + grpc_server.wait_for_termination() + except KeyboardInterrupt: + server_instance.stop() + grpc_server.stop(grace=2) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8080) + parser.add_argument("--planner_url", default="http://localhost:9090/plan") + parser.add_argument("--environment_dt", type=float, default=1.0 / 30) + parser.add_argument("--obs_queue_timeout", type=float, default=1.0) + parser.add_argument("--include_images", action="store_true") + parser.add_argument("--http_timeout", type=float, default=5.0) + args = parser.parse_args() + + serve( + host=args.host, + port=args.port, + planner_url=args.planner_url, + environment_dt=args.environment_dt, + obs_queue_timeout=args.obs_queue_timeout, + include_images=args.include_images, + http_timeout=args.http_timeout, + ) diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/hardcoded_action_server.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/hardcoded_action_server.py new file mode 100644 index 0000000..1741f81 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/hardcoded_action_server.py @@ -0,0 +1,219 @@ +""" +Minimal gRPC server that returns hardcoded actions derived from the robot's own +current joint positions. No external planner, no neural network. + +Purpose: verify that the full pipeline (RobotClient -> gRPC -> robot.send_action) +actually moves the robot before connecting a real planner. + +HOW IT WORKS +------------ +1. Robot sends its current joint positions as an observation. +2. This server reads those positions and builds a small action chunk: + step 0 : current position (hold still) + step 1 : current position + DELTA (tiny movement) + step 2 : current position (return) +3. RobotClient executes those actions in order -> robot should wiggle slightly. + +USAGE +----- +Terminal 1 (this server): + python -m lerobot.async_inference.hardcoded_action_server \ + --port=8080 \ + --delta=2.0 \ + --joint_index=0 + +Terminal 2 (real robot): + python -m lerobot.async_inference.robot_client \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyUSB0 \ + --robot.id=my_robot \ + --server_address=127.0.0.1:8080 \ + --task="test" \ + --policy_type=act \ + --pretrained_name_or_path=test/model \ + --actions_per_chunk=3 \ + --policy_device=cpu \ + --client_device=cpu + +SAFETY +------ +- DELTA default is 2.0 degrees. Start small. +- The robot moves to current+DELTA then returns. Very short wiggle. +- Stop anytime with Ctrl+C on Terminal 2. +""" + +import argparse +import logging +import pickle # nosec +import threading +import time +from concurrent import futures +from queue import Empty, Queue + +import grpc +import numpy as np +import torch + +from lerobot.transport import services_pb2, services_pb2_grpc +from lerobot.transport.utils import receive_bytes_in_chunks + +from .helpers import TimedAction, TimedObservation + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger("hardcoded_action_server") + + +class HardcodedActionServer(services_pb2_grpc.AsyncInferenceServicer): + """ + Returns actions built from the robot's own current state. + Proves the data pipeline works without any external planner. + """ + + def __init__(self, delta: float = 2.0, joint_index: int = 0, environment_dt: float = 1 / 30): + """ + Args: + delta: How many degrees to move joint `joint_index` (default 2.0). + joint_index: Which joint to wiggle (0 = first joint in action_features order). + environment_dt: Control period in seconds. + """ + self.delta = delta + self.joint_index = joint_index + self.environment_dt = environment_dt + + self.shutdown_event = threading.Event() + self.observation_queue: Queue[TimedObservation] = Queue(maxsize=1) + + @property + def running(self): + return not self.shutdown_event.is_set() + + # ------------------------------------------------------------------ + # gRPC handlers + # ------------------------------------------------------------------ + + def Ready(self, request, context): # noqa: N802 + logger.info(f"Robot client connected: {context.peer()}") + self.shutdown_event.clear() + self.observation_queue = Queue(maxsize=1) + return services_pb2.Empty() + + def SendPolicyInstructions(self, request, context): # noqa: N802 + policy_specs = pickle.loads(request.data) # nosec + logger.info( + f"Policy instructions received | " + f"policy_type={policy_specs.policy_type} | " + f"actions_per_chunk={policy_specs.actions_per_chunk}" + ) + return services_pb2.Empty() + + def SendObservations(self, request_iterator, context): # noqa: N802 + received_bytes = receive_bytes_in_chunks( + request_iterator, None, self.shutdown_event, logger + ) + timed_obs: TimedObservation = pickle.loads(received_bytes) # nosec + + raw_obs = timed_obs.get_observation() + + # Raw observation uses motor names as keys (e.g. "shoulder_pan.pos"), + # not the dataset key "observation.state". + motor_keys = [k for k in raw_obs if k.endswith(".pos")] + logger.info( + f"Obs #{timed_obs.get_timestep()} | " + f"must_go={timed_obs.must_go} | " + f"motor_keys={motor_keys} | " + f"values={[round(float(raw_obs[k]), 2) for k in motor_keys]}" + ) + + if self.observation_queue.full(): + try: + self.observation_queue.get_nowait() + except Empty: + pass + self.observation_queue.put(timed_obs) + return services_pb2.Empty() + + def GetActions(self, request, context): # noqa: N802 + try: + obs: TimedObservation = self.observation_queue.get(timeout=1.0) + except Empty: + logger.debug("No observation in queue, returning empty.") + return services_pb2.Actions(data=b"") + + raw_obs = obs.get_observation() + + # Collect motor positions in sorted key order (must match robot.action_features order) + motor_keys = [k for k in raw_obs if k.endswith(".pos")] + + if not motor_keys: + logger.warning(f"No '.pos' keys found in observation. Keys present: {list(raw_obs.keys())}") + return services_pb2.Actions(data=b"") + + current = np.array([float(raw_obs[k]) for k in motor_keys], dtype=np.float32) + action_dim = len(current) + + # Build 3-step chunk: + # step 0: hold current position + # step 1: move joint[joint_index] by +delta + # step 2: return to current position + step0 = current.copy() + step1 = current.copy() + step1[self.joint_index % action_dim] += self.delta + step2 = current.copy() + + timed_actions = [ + TimedAction( + timestamp=obs.get_timestamp() + i * self.environment_dt, + timestep=obs.get_timestep() + i, + action=torch.tensor(a, dtype=torch.float32), + ) + for i, a in enumerate([step0, step1, step2]) + ] + + logger.info( + f"Sending 3 actions for obs #{obs.get_timestep()} | " + f"joint[{self.joint_index}]: " + f"{current[self.joint_index % action_dim]:.2f} -> " + f"{step1[self.joint_index % action_dim]:.2f} -> " + f"{step2[self.joint_index % action_dim]:.2f}" + ) + + return services_pb2.Actions(data=pickle.dumps(timed_actions)) # nosec + + def stop(self): + self.shutdown_event.set() + logger.info("Server stopped.") + + +def serve(port: int = 8080, delta: float = 2.0, joint_index: int = 0, fps: int = 30): + server_instance = HardcodedActionServer( + delta=delta, + joint_index=joint_index, + environment_dt=1.0 / fps, + ) + + grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + services_pb2_grpc.add_AsyncInferenceServicer_to_server(server_instance, grpc_server) + grpc_server.add_insecure_port(f"0.0.0.0:{port}") + + logger.info(f"HardcodedActionServer listening on port {port}") + logger.info(f"Will wiggle joint[{joint_index}] by ยฑ{delta} degrees each observation") + grpc_server.start() + + try: + grpc_server.wait_for_termination() + except KeyboardInterrupt: + server_instance.stop() + grpc_server.stop(grace=2) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--port", type=int, default=8080) + parser.add_argument("--delta", type=float, default=2.0, + help="Degrees to move the test joint (keep small!)") + parser.add_argument("--joint_index", type=int, default=0, + help="Which joint to wiggle (0-indexed)") + parser.add_argument("--fps", type=int, default=30) + args = parser.parse_args() + + serve(port=args.port, delta=args.delta, joint_index=args.joint_index, fps=args.fps) diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/helpers.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/helpers.py new file mode 100644 index 0000000..c8cadd3 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/helpers.py @@ -0,0 +1,297 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import logging.handlers +import os +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch + +from lerobot.configs.types import PolicyFeature +from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features + +# NOTE: Configs need to be loaded for the client to be able to instantiate the policy config +from lerobot.policies import ( # noqa: F401 + ACTConfig, + DiffusionConfig, + PI0Config, + PI05Config, + SmolVLAConfig, + VQBeTConfig, +) +from lerobot.robots.robot import Robot +from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR +from lerobot.utils.utils import init_logging + +Action = torch.Tensor + +# observation as received from the robot (can be numpy arrays, floats, etc.) +RawObservation = dict[str, Any] + +# observation as those recorded in LeRobot dataset (keys are different) +LeRobotObservation = dict[str, torch.Tensor] + +# observation, ready for policy inference (image keys resized) +Observation = dict[str, torch.Tensor] + + +def visualize_action_queue_size(action_queue_size: list[int]) -> None: + import matplotlib.pyplot as plt + + _, ax = plt.subplots() + ax.set_title("Action Queue Size Over Time") + ax.set_xlabel("Environment steps") + ax.set_ylabel("Action Queue Size") + ax.set_ylim(0, max(action_queue_size) * 1.1) + ax.grid(True, alpha=0.3) + ax.plot(range(len(action_queue_size)), action_queue_size) + plt.show() + + +def map_robot_keys_to_lerobot_features(robot: Robot) -> dict[str, dict]: + return hw_to_dataset_features(robot.observation_features, OBS_STR, use_video=False) + + +def is_image_key(k: str) -> bool: + return k.startswith(OBS_IMAGES) + + +def resize_robot_observation_image(image: torch.tensor, resize_dims: tuple[int, int, int]) -> torch.tensor: + assert image.ndim == 3, f"Image must be (C, H, W)! Received {image.shape}" + # (H, W, C) -> (C, H, W) for resizing from robot obsevation resolution to policy image resolution + image = image.permute(2, 0, 1) + dims = (resize_dims[1], resize_dims[2]) + # Add batch dimension for interpolate: (C, H, W) -> (1, C, H, W) + image_batched = image.unsqueeze(0) + # Interpolate and remove batch dimension: (1, C, H, W) -> (C, H, W) + resized = torch.nn.functional.interpolate(image_batched, size=dims, mode="bilinear", align_corners=False) + + return resized.squeeze(0) + + +# TODO(Steven): Consider implementing a pipeline step for this +def raw_observation_to_observation( + raw_observation: RawObservation, + lerobot_features: dict[str, dict], + policy_image_features: dict[str, PolicyFeature], +) -> Observation: + observation = {} + + observation = prepare_raw_observation(raw_observation, lerobot_features, policy_image_features) + for k, v in observation.items(): + if isinstance(v, torch.Tensor): # VLAs present natural-language instructions in observations + if "image" in k: + # Policy expects images in shape (B, C, H, W) + observation[k] = prepare_image(v).unsqueeze(0) + else: + observation[k] = v + + return observation + + +def prepare_image(image: torch.Tensor) -> torch.Tensor: + """Minimal preprocessing to turn int8 images to float32 in [0, 1], and create a memory-contiguous tensor""" + image = image.type(torch.float32) / 255 + image = image.contiguous() + + return image + + +def extract_state_from_raw_observation( + lerobot_obs: RawObservation, +) -> torch.Tensor: + """Extract the state from a raw observation.""" + state = torch.tensor(lerobot_obs[OBS_STATE]) + + if state.ndim == 1: + state = state.unsqueeze(0) + + return state + + +def extract_images_from_raw_observation( + lerobot_obs: RawObservation, + camera_key: str, +) -> dict[str, torch.Tensor]: + """Extract the images from a raw observation.""" + return torch.tensor(lerobot_obs[camera_key]) + + +def make_lerobot_observation( + robot_obs: RawObservation, + lerobot_features: dict[str, dict], +) -> LeRobotObservation: + """Make a lerobot observation from a raw observation.""" + return build_dataset_frame(lerobot_features, robot_obs, prefix=OBS_STR) + + +def prepare_raw_observation( + robot_obs: RawObservation, + lerobot_features: dict[str, dict], + policy_image_features: dict[str, PolicyFeature], +) -> Observation: + """Matches keys from the raw robot_obs dict to the keys expected by a given policy (passed as + policy_image_features).""" + # 1. {motor.pos1:value1, motor.pos2:value2, ..., laptop:np.ndarray} -> + # -> {observation.state:[value1,value2,...], observation.images.laptop:np.ndarray} + lerobot_obs = make_lerobot_observation(robot_obs, lerobot_features) + + # 2. Greps all observation.images.<> keys + image_keys = list(filter(is_image_key, lerobot_obs)) + # state's shape is expected as (B, state_dim) + state_dict = {OBS_STATE: extract_state_from_raw_observation(lerobot_obs)} + image_dict = { + image_k: extract_images_from_raw_observation(lerobot_obs, image_k) for image_k in image_keys + } + + # Turns the image features to (C, H, W) with H, W matching the policy image features. + # This reduces the resolution of the images + image_dict = { + key: resize_robot_observation_image(torch.tensor(lerobot_obs[key]), policy_image_features[key].shape) + for key in image_keys + } + + if "task" in robot_obs: + state_dict["task"] = robot_obs["task"] + + return {**state_dict, **image_dict} + + +def get_logger(name: str, log_to_file: bool = True) -> logging.Logger: + """ + Get a logger using the standardized logging setup from utils.py. + + Args: + name: Logger name (e.g., 'policy_server', 'robot_client') + log_to_file: Whether to also log to a file + + Returns: + Configured logger instance + """ + # Create logs directory if logging to file + if log_to_file: + os.makedirs("logs", exist_ok=True) + log_file = Path(f"logs/{name}_{int(time.time())}.log") + else: + log_file = None + + # Initialize the standardized logging + init_logging(log_file=log_file, display_pid=False) + + # Return a named logger + return logging.getLogger(name) + + +@dataclass +class TimedData: + """A data object with timestamp and timestep information. + + Args: + timestamp: Unix timestamp relative to data's creation. + data: The actual data to wrap a timestamp around. + timestep: The timestep of the data. + """ + + timestamp: float + timestep: int + + def get_timestamp(self): + return self.timestamp + + def get_timestep(self): + return self.timestep + + +@dataclass +class TimedAction(TimedData): + action: Action + + def get_action(self): + return self.action + + +@dataclass +class TimedObservation(TimedData): + observation: RawObservation + must_go: bool = False + + def get_observation(self): + return self.observation + + +@dataclass +class FPSTracker: + """Utility class to track FPS metrics over time.""" + + target_fps: float + first_timestamp: float = None + total_obs_count: int = 0 + + def calculate_fps_metrics(self, current_timestamp: float) -> dict[str, float]: + """Calculate average FPS vs target""" + self.total_obs_count += 1 + + # Initialize first observation time + if self.first_timestamp is None: + self.first_timestamp = current_timestamp + + # Calculate overall average FPS (since start) + total_duration = current_timestamp - self.first_timestamp + avg_fps = (self.total_obs_count - 1) / total_duration if total_duration > 1e-6 else 0.0 + + return {"avg_fps": avg_fps, "target_fps": self.target_fps} + + def reset(self): + """Reset the FPS tracker state""" + self.first_timestamp = None + self.total_obs_count = 0 + + +@dataclass +class RemotePolicyConfig: + policy_type: str + pretrained_name_or_path: str + lerobot_features: dict[str, PolicyFeature] + actions_per_chunk: int + device: str = "cpu" + rename_map: dict[str, str] = field(default_factory=dict) + + +def _compare_observation_states(obs1_state: torch.Tensor, obs2_state: torch.Tensor, atol: float) -> bool: + """Check if two observation states are similar, under a tolerance threshold""" + return bool(torch.linalg.norm(obs1_state - obs2_state) < atol) + + +def observations_similar( + obs1: TimedObservation, obs2: TimedObservation, lerobot_features: dict[str, dict], atol: float = 1 +) -> bool: + """Check if two observations are similar, under a tolerance threshold. Measures distance between + observations as the difference in joint-space between the two observations. + + NOTE(fracapuano): This is a very simple check, and it is enough for the current use case. + An immediate next step is to use (fast) perceptual difference metrics comparing some camera views, + to surpass this joint-space similarity check. + """ + obs1_state = extract_state_from_raw_observation( + make_lerobot_observation(obs1.get_observation(), lerobot_features) + ) + obs2_state = extract_state_from_raw_observation( + make_lerobot_observation(obs2.get_observation(), lerobot_features) + ) + + return _compare_observation_states(obs1_state, obs2_state, atol=atol) diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/policy_server.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/policy_server.py new file mode 100644 index 0000000..3f63929 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/policy_server.py @@ -0,0 +1,439 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Example: +```shell +python -m lerobot.async_inference.policy_server \ + --host=127.0.0.1 \ + --port=8080 \ + --fps=30 \ + --inference_latency=0.033 \ + --obs_queue_timeout=1 +``` +""" + +import logging +import pickle # nosec +import threading +import time +from concurrent import futures +from dataclasses import asdict +from pprint import pformat +from queue import Empty, Queue +from typing import Any + +import draccus +import grpc +import torch + +from lerobot.policies.factory import get_policy_class, make_pre_post_processors +from lerobot.processor import PolicyProcessorPipeline +from lerobot.transport import ( + services_pb2, # type: ignore + services_pb2_grpc, # type: ignore +) +from lerobot.transport.utils import receive_bytes_in_chunks +from lerobot.types import PolicyAction + +from .configs import PolicyServerConfig +from .constants import SUPPORTED_POLICIES +from .helpers import ( + FPSTracker, + Observation, + RemotePolicyConfig, + TimedAction, + TimedObservation, + get_logger, + observations_similar, + raw_observation_to_observation, +) + + +class PolicyServer(services_pb2_grpc.AsyncInferenceServicer): + prefix = "policy_server" + logger = get_logger(prefix) + + def __init__(self, config: PolicyServerConfig): + self.config = config + self.shutdown_event = threading.Event() + + # FPS measurement + self.fps_tracker = FPSTracker(target_fps=config.fps) + + self.observation_queue = Queue(maxsize=1) + + self._predicted_timesteps_lock = threading.Lock() + self._predicted_timesteps = set() + + self.last_processed_obs = None + + # Attributes will be set by SendPolicyInstructions + self.device = None + self.policy_type = None + self.lerobot_features = None + self.actions_per_chunk = None + self.policy = None + self.preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None + self.postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None + + @property + def running(self): + return not self.shutdown_event.is_set() + + @property + def policy_image_features(self): + return self.policy.config.image_features + + def _reset_server(self) -> None: + """Flushes server state when new client connects.""" + # only running inference on the latest observation received by the server + self.shutdown_event.set() + self.observation_queue = Queue(maxsize=1) + + with self._predicted_timesteps_lock: + self._predicted_timesteps = set() + + def Ready(self, request, context): # noqa: N802 + client_id = context.peer() + self.logger.info(f"Client {client_id} connected and ready") + self._reset_server() + self.shutdown_event.clear() + + return services_pb2.Empty() + + def SendPolicyInstructions(self, request, context): # noqa: N802 + """Receive policy instructions from the robot client""" + + if not self.running: + self.logger.warning("Server is not running. Ignoring policy instructions.") + return services_pb2.Empty() + + client_id = context.peer() + + policy_specs = pickle.loads(request.data) # nosec + + if not isinstance(policy_specs, RemotePolicyConfig): + raise TypeError(f"Policy specs must be a RemotePolicyConfig. Got {type(policy_specs)}") + + if policy_specs.policy_type not in SUPPORTED_POLICIES: + raise ValueError( + f"Policy type {policy_specs.policy_type} not supported. " + f"Supported policies: {SUPPORTED_POLICIES}" + ) + + self.logger.info( + f"Receiving policy instructions from {client_id} | " + f"Policy type: {policy_specs.policy_type} | " + f"Pretrained name or path: {policy_specs.pretrained_name_or_path} | " + f"Actions per chunk: {policy_specs.actions_per_chunk} | " + f"Device: {policy_specs.device}" + ) + + self.device = policy_specs.device + self.policy_type = policy_specs.policy_type # act, pi0, etc. + self.lerobot_features = policy_specs.lerobot_features + self.actions_per_chunk = policy_specs.actions_per_chunk + + policy_class = get_policy_class(self.policy_type) + + start = time.perf_counter() + self.policy = policy_class.from_pretrained(policy_specs.pretrained_name_or_path) + self.policy.to(self.device) + + # Load preprocessor and postprocessor, overriding device to match requested device + device_override = {"device": self.device} + self.preprocessor, self.postprocessor = make_pre_post_processors( + self.policy.config, + pretrained_path=policy_specs.pretrained_name_or_path, + preprocessor_overrides={ + "device_processor": device_override, + "rename_observations_processor": {"rename_map": policy_specs.rename_map}, + }, + postprocessor_overrides={"device_processor": device_override}, + ) + + end = time.perf_counter() + + self.logger.info(f"Time taken to put policy on {self.device}: {end - start:.4f} seconds") + + return services_pb2.Empty() + + def SendObservations(self, request_iterator, context): # noqa: N802 + """Receive observations from the robot client""" + client_id = context.peer() + self.logger.debug(f"Receiving observations from {client_id}") + + receive_time = time.time() # comparing timestamps so need time.time() + start_deserialize = time.perf_counter() + received_bytes = receive_bytes_in_chunks( + request_iterator, None, self.shutdown_event, self.logger + ) # blocking call while looping over request_iterator + timed_observation = pickle.loads(received_bytes) # nosec + deserialize_time = time.perf_counter() - start_deserialize + + self.logger.debug(f"Received observation #{timed_observation.get_timestep()}") + + obs_timestep = timed_observation.get_timestep() + obs_timestamp = timed_observation.get_timestamp() + + # Calculate FPS metrics + fps_metrics = self.fps_tracker.calculate_fps_metrics(obs_timestamp) + + self.logger.debug( + f"Received observation #{obs_timestep} | " + f"Avg FPS: {fps_metrics['avg_fps']:.2f} | " # fps at which observations are received from client + f"Target: {fps_metrics['target_fps']:.2f} | " + f"One-way latency: {(receive_time - obs_timestamp) * 1000:.2f}ms" + ) + + self.logger.debug( + f"Server timestamp: {receive_time:.6f} | " + f"Client timestamp: {obs_timestamp:.6f} | " + f"Deserialization time: {deserialize_time:.6f}s" + ) + + if not self._enqueue_observation( + timed_observation # wrapping a RawObservation + ): + self.logger.debug(f"Observation #{obs_timestep} has been filtered out") + + return services_pb2.Empty() + + def GetActions(self, request, context): # noqa: N802 + """Returns actions to the robot client. Actions are sent as a single + chunk, containing multiple actions.""" + client_id = context.peer() + self.logger.debug(f"Client {client_id} connected for action streaming") + + # Generate action based on the most recent observation and its timestep + try: + getactions_starts = time.perf_counter() + obs = self.observation_queue.get(timeout=self.config.obs_queue_timeout) + self.logger.info( + f"Running inference for observation #{obs.get_timestep()} (must_go: {obs.must_go})" + ) + + with self._predicted_timesteps_lock: + self._predicted_timesteps.add(obs.get_timestep()) + + start_time = time.perf_counter() + action_chunk = self._predict_action_chunk(obs) + inference_time = time.perf_counter() - start_time + + start_time = time.perf_counter() + actions_bytes = pickle.dumps(action_chunk) # nosec + serialize_time = time.perf_counter() - start_time + + # Create and return the action chunk + actions = services_pb2.Actions(data=actions_bytes) + + self.logger.info( + f"Action chunk #{obs.get_timestep()} generated | " + f"Total time: {(inference_time + serialize_time) * 1000:.2f}ms" + ) + + self.logger.debug( + f"Action chunk #{obs.get_timestep()} generated | " + f"Inference time: {inference_time:.2f}s |" + f"Serialize time: {serialize_time:.2f}s |" + f"Total time: {inference_time + serialize_time:.2f}s" + ) + + time.sleep( + max(0, self.config.inference_latency - max(0, time.perf_counter() - getactions_starts)) + ) # sleep controls inference latency + + return actions + + except Empty: # no observation added to queue in obs_queue_timeout + return services_pb2.Empty() + + except Exception as e: + self.logger.error(f"Error in StreamActions: {e}") + + return services_pb2.Empty() + + def _obs_sanity_checks(self, obs: TimedObservation, previous_obs: TimedObservation) -> bool: + """Check if the observation is valid to be processed by the policy""" + with self._predicted_timesteps_lock: + predicted_timesteps = self._predicted_timesteps + + if obs.get_timestep() in predicted_timesteps: + self.logger.debug(f"Skipping observation #{obs.get_timestep()} - Timestep predicted already!") + return False + + elif observations_similar(obs, previous_obs, lerobot_features=self.lerobot_features): + self.logger.debug( + f"Skipping observation #{obs.get_timestep()} - Observation too similar to last obs predicted!" + ) + return False + + else: + return True + + def _enqueue_observation(self, obs: TimedObservation) -> bool: + """Enqueue an observation if it must go through processing, otherwise skip it. + Observations not in queue are never run through the policy network""" + + if ( + obs.must_go + or self.last_processed_obs is None + or self._obs_sanity_checks(obs, self.last_processed_obs) + ): + last_obs = self.last_processed_obs.get_timestep() if self.last_processed_obs else "None" + self.logger.debug( + f"Enqueuing observation. Must go: {obs.must_go} | Last processed obs: {last_obs}" + ) + + # If queue is full, get the old observation to make room + if self.observation_queue.full(): + # pops from queue + _ = self.observation_queue.get_nowait() + self.logger.debug("Observation queue was full, removed oldest observation") + + # Now put the new observation (never blocks as queue is non-full here) + self.observation_queue.put(obs) + return True + + return False + + def _time_action_chunk(self, t_0: float, action_chunk: list[torch.Tensor], i_0: int) -> list[TimedAction]: + """Turn a chunk of actions into a list of TimedAction instances, + with the first action corresponding to t_0 and the rest corresponding to + t_0 + i*environment_dt for i in range(len(action_chunk)) + """ + return [ + TimedAction(timestamp=t_0 + i * self.config.environment_dt, timestep=i_0 + i, action=action) + for i, action in enumerate(action_chunk) + ] + + def _get_action_chunk(self, observation: dict[str, torch.Tensor]) -> torch.Tensor: + """Get an action chunk from the policy. The chunk contains only""" + chunk = self.policy.predict_action_chunk(observation) + if chunk.ndim != 3: + chunk = chunk.unsqueeze(0) # adding batch dimension, now shape is (B, chunk_size, action_dim) + + return chunk[:, : self.actions_per_chunk, :] + + def _predict_action_chunk(self, observation_t: TimedObservation) -> list[TimedAction]: + """Predict an action chunk based on an observation. + + Pipeline: + 1. Convert raw observation to LeRobot format + 2. Apply preprocessor (tokenization, normalization, batching, device placement) + 3. Run policy inference to get action chunk + 4. Apply postprocessor (unnormalization, device movement) + 5. Convert to TimedAction list + """ + """1. Prepare observation""" + start_prepare = time.perf_counter() + observation: Observation = raw_observation_to_observation( + observation_t.get_observation(), + self.lerobot_features, + self.policy_image_features, + ) + prepare_time = time.perf_counter() - start_prepare + + """2. Apply preprocessor""" + start_preprocess = time.perf_counter() + observation = self.preprocessor(observation) + self.last_processed_obs: TimedObservation = observation_t + preprocessing_time = time.perf_counter() - start_preprocess + + """3. Get action chunk""" + start_inference = time.perf_counter() + action_tensor = self._get_action_chunk(observation) + inference_time = time.perf_counter() - start_inference + self.logger.info( + f"Preprocessing and inference took {inference_time:.4f}s, action shape: {action_tensor.shape}" + ) + + """4. Apply postprocessor""" + # Apply postprocessor (handles unnormalization and device movement) + # Postprocessor expects (B, action_dim) per action, but we have (B, chunk_size, action_dim) + # So we process each action in the chunk individually + start_postprocess = time.perf_counter() + _, chunk_size, _ = action_tensor.shape + + # Process each action in the chunk + processed_actions = [] + for i in range(chunk_size): + # Extract action at timestep i: (B, action_dim) + single_action = action_tensor[:, i, :] + processed_action = self.postprocessor(single_action) + processed_actions.append(processed_action) + + # Stack back to (B, chunk_size, action_dim), then remove batch dim + action_tensor = torch.stack(processed_actions, dim=1).squeeze(0) + self.logger.debug(f"Postprocessed action shape: {action_tensor.shape}") + + action_tensor = action_tensor.detach().cpu() + + """5. Convert to TimedAction list""" + action_chunk = self._time_action_chunk( + observation_t.get_timestamp(), list(action_tensor), observation_t.get_timestep() + ) + postprocess_stops = time.perf_counter() + postprocessing_time = postprocess_stops - start_postprocess + + self.logger.info( + f"Observation {observation_t.get_timestep()} | " + f"Total time: {1000 * (postprocess_stops - start_prepare):.2f}ms" + ) + + self.logger.debug( + f"Observation {observation_t.get_timestep()} | " + f"Prepare time: {1000 * prepare_time:.2f}ms | " + f"Preprocessing time: {1000 * preprocessing_time:.2f}ms | " + f"Inference time: {1000 * inference_time:.2f}ms | " + f"Postprocessing time: {1000 * postprocessing_time:.2f}ms | " + f"Total time: {1000 * (postprocess_stops - start_prepare):.2f}ms" + ) + + return action_chunk + + def stop(self): + """Stop the server""" + self._reset_server() + self.logger.info("Server stopping...") + + +@draccus.wrap() +def serve(cfg: PolicyServerConfig): + """Start the PolicyServer with the given configuration. + + Args: + config: PolicyServerConfig instance. If None, uses default configuration. + """ + logging.info(pformat(asdict(cfg))) + + # Create the server instance first + policy_server = PolicyServer(cfg) + + # Setup and start gRPC server + server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + services_pb2_grpc.add_AsyncInferenceServicer_to_server(policy_server, server) + server.add_insecure_port(f"{cfg.host}:{cfg.port}") + + policy_server.logger.info(f"PolicyServer started on {cfg.host}:{cfg.port}") + server.start() + + server.wait_for_termination() + + policy_server.logger.info("Server terminated") + + +if __name__ == "__main__": + serve() diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/rl_client.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/rl_client.py new file mode 100644 index 0000000..1600bc7 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/rl_client.py @@ -0,0 +1,212 @@ +""" +LeRobot RL client โ€” runs the episode loop between the robot server and the policy server. + +Data flow per timestep: + 1. GET robot_server:7070/state โ†’ obs (6 joint floats) + 2. POST policy_server:8765/act โ† {"obs": [6 floats]} + โ†’ {"action": [6 floats]} + 3. POST robot_server:7070/step โ† {"joints": [6 floats]} + โ†’ next_obs + 4. POST policy_server:8765/step_done โ† {"obs": [6 floats], + "action": [6 floats], + "next_obs": [6 floats], + "done": bool} + +At episode end: + 5. POST robot_server:7070/reset โ†’ go to home position + 6. POST policy_server:8765/reset โ†’ triggers gradient update, clears storage + +Usage: + python -m lerobot.async_inference.rl_client \\ + --robot_url=http://127.0.0.1:7070 \\ + --policy_url=http://127.0.0.1:8765 \\ + --num_episodes=20 \\ + --max_steps=200 \\ + --control_hz=10.0 +""" + +import logging +import time +from dataclasses import dataclass, field + +import draccus +import requests + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger("rl_client") + + +@dataclass +class RLClientConfig: + robot_url: str = "http://127.0.0.1:7070" + policy_url: str = "http://127.0.0.1:8765" + num_episodes: int = 20 + max_steps: int = 200 + control_hz: float = 10.0 + # Seconds before giving up on an HTTP call + timeout: float = 5.0 + # Home position (degrees) sent to the robot on every episode reset. + # Override with --home_joints='[16.44,-104.44,95.91,72.75,-0.04,1.96]' + home_joints: list = field( + default_factory=lambda: [16.44, -104.44, 95.91, 72.75, -0.04, 1.96] + ) + + +# --------------------------------------------------------------------------- +# Thin HTTP helpers +# --------------------------------------------------------------------------- + +def _get(url: str, timeout: float) -> dict: + r = requests.get(url, timeout=timeout) + r.raise_for_status() + return r.json() + + +def _post(url: str, body: dict, timeout: float) -> dict: + r = requests.post(url, json=body, timeout=timeout) + r.raise_for_status() + return r.json() + + +# --------------------------------------------------------------------------- +# Per-server calls +# --------------------------------------------------------------------------- + +def robot_get_state(robot_url: str, timeout: float) -> list[float]: + data = _get(f"{robot_url}/state", timeout) + return data["joints"] + + +def robot_step(robot_url: str, joints: list[float], timeout: float) -> list[float]: + data = _post(f"{robot_url}/step", {"joints": joints}, timeout) + return data["joints"] + + +def robot_reset(robot_url: str, timeout: float, home_joints: list[float] | None = None) -> list[float]: + body = {"joints": home_joints} if home_joints else {} + data = _post(f"{robot_url}/reset", body, timeout) + return data["joints"] + + +def policy_act(policy_url: str, obs: list[float], timeout: float) -> list[float]: + data = _post(f"{policy_url}/act", {"obs": obs}, timeout) + return data["action"] + + +def policy_step_done( + policy_url: str, + obs: list[float], + action: list[float], + next_obs: list[float], + done: bool, + timeout: float, + reward: float = 0.0, +) -> None: + # reward=0.0 โ€” MPAIL learns reward from demonstrations, not from the env signal + _post( + f"{policy_url}/step_done", + {"obs": obs, "action": action, "next_obs": next_obs, "done": done, "reward": reward}, + timeout, + ) + + +def policy_reset(policy_url: str, timeout: float) -> None: + _post(f"{policy_url}/reset", {}, timeout) + + +# --------------------------------------------------------------------------- +# Episode loop +# --------------------------------------------------------------------------- + +def run_episode(cfg: RLClientConfig, episode: int) -> int: + """Run one episode. Returns the number of steps taken.""" + dt = 1.0 / cfg.control_hz + + # Start from a clean robot state + obs = robot_reset(cfg.robot_url, cfg.timeout, cfg.home_joints) + logger.info(f"Episode {episode} start | obs={[round(v,2) for v in obs]}") + + for step in range(cfg.max_steps): + step_start = time.perf_counter() + + # 1. Get action from policy server + try: + action = policy_act(cfg.policy_url, obs, cfg.timeout) + except requests.RequestException as e: + logger.error(f" /act failed at step {step}: {e}") + break + + # 2. Send action to robot, get next observation + try: + next_obs = robot_step(cfg.robot_url, action, cfg.timeout) + except requests.RequestException as e: + logger.error(f" /step failed at step {step}: {e}") + break + + done = (step == cfg.max_steps - 1) + + # 3. Report transition to policy server + try: + policy_step_done(cfg.policy_url, obs, action, next_obs, done, cfg.timeout) + except requests.RequestException as e: + logger.warning(f" /step_done failed at step {step}: {e}") + + logger.debug( + f" step {step:3d} | " + f"action={[round(a,1) for a in action]} | " + f"next_obs={[round(v,2) for v in next_obs]}" + ) + + obs = next_obs + + if done: + break + + # Pace the loop to control_hz + elapsed = time.perf_counter() - step_start + time.sleep(max(0.0, dt - elapsed)) + + steps_taken = step + 1 + logger.info(f"Episode {episode} done after {steps_taken} steps โ€” sending /reset to policy server") + + # 4. Reset policy server (triggers gradient update) + try: + policy_reset(cfg.policy_url, cfg.timeout) + except requests.RequestException as e: + logger.warning(f" /reset (policy) failed: {e}") + + return steps_taken + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +@draccus.wrap() +def main(cfg: RLClientConfig): + # Health check + try: + _get(f"{cfg.policy_url}/health", timeout=3.0) + logger.info(f"Policy server reachable at {cfg.policy_url}") + except requests.RequestException as e: + logger.error(f"Policy server not reachable at {cfg.policy_url}: {e}") + return + + try: + robot_get_state(cfg.robot_url, timeout=3.0) + logger.info(f"Robot server reachable at {cfg.robot_url} home={cfg.home_joints}") + except requests.RequestException as e: + logger.error(f"Robot server not reachable at {cfg.robot_url}: {e}") + return + + total_steps = 0 + for ep in range(1, cfg.num_episodes + 1): + steps = run_episode(cfg, ep) + total_steps += steps + logger.info(f"Total steps so far: {total_steps}") + + logger.info(f"Done โ€” {cfg.num_episodes} episodes, {total_steps} total steps.") + + +if __name__ == "__main__": + main() diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/robot_client.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/robot_client.py new file mode 100644 index 0000000..bd01080 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/robot_client.py @@ -0,0 +1,738 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Example command: +```shell +python src/lerobot/async_inference/robot_client.py \ + --robot.type=so100_follower \ + --robot.port=/dev/tty.usbmodem58760431541 \ + --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \ + --robot.id=black \ + --task="dummy" \ + --server_address=127.0.0.1:8080 \ + --policy_type=act \ + --pretrained_name_or_path=user/model \ + --policy_device=mps \ + --client_device=cpu \ + --actions_per_chunk=50 \ + --chunk_size_threshold=0.5 \ + --aggregate_fn_name=weighted_average \ + --debug_visualize_queue_size=True +``` +""" + +import logging +import pickle # nosec +import threading +import time +from collections.abc import Callable +from dataclasses import asdict +from pprint import pformat +from queue import Queue +from typing import Any + +import draccus +import grpc +import torch + +from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.robots import ( # noqa: F401 + Robot, + RobotConfig, + bi_so_follower, + koch_follower, + make_robot_from_config, + omx_follower, + so_follower, +) +from lerobot.teleoperators import ( # noqa: F401 + bi_so_leader, + koch_leader, + make_teleoperator_from_config, + omx_leader, + so_leader, +) +from lerobot.transport import ( + services_pb2, # type: ignore + services_pb2_grpc, # type: ignore +) +from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks +from lerobot.utils.import_utils import register_third_party_plugins + +from .configs import RobotClientConfig +from .helpers import ( + Action, + FPSTracker, + Observation, + RawObservation, + RemotePolicyConfig, + TimedAction, + TimedObservation, + get_logger, + map_robot_keys_to_lerobot_features, + visualize_action_queue_size, +) + + +class RobotClient: + prefix = "robot_client" + logger = get_logger(prefix) + + def __init__(self, config: RobotClientConfig): + """Initialize RobotClient with unified configuration. + + Args: + config: RobotClientConfig containing all configuration parameters + """ + # Store configuration + self.config = config + self.robot = make_robot_from_config(config.robot) + self.robot.connect() + + self.teleop = make_teleoperator_from_config(config.teleop) if config.teleop else None + if self.teleop: + self.teleop.connect() + self.logger.info("Teleoperator (leader arm) connected โ€” running in teleoperation mode.") + + lerobot_features = map_robot_keys_to_lerobot_features(self.robot) + + # Use environment variable if server_address is not provided in config + self.server_address = config.server_address + + self.policy_config = RemotePolicyConfig( + config.policy_type, + config.pretrained_name_or_path, + lerobot_features, + config.actions_per_chunk, + config.policy_device, + ) + self.channel = grpc.insecure_channel( + self.server_address, grpc_channel_options(initial_backoff=f"{config.environment_dt:.4f}s") + ) + self.stub = services_pb2_grpc.AsyncInferenceStub(self.channel) + self.logger.info(f"Initializing client to connect to server at {self.server_address}") + + self.shutdown_event = threading.Event() + + # Initialize client side variables + self.latest_action_lock = threading.Lock() + self.latest_action = -1 + self.action_chunk_size = -1 + + self._chunk_size_threshold = config.chunk_size_threshold + + self.action_queue = Queue() + self.action_queue_lock = threading.Lock() # Protect queue operations + self.action_queue_size = [] + self.start_barrier = threading.Barrier(2) # 2 threads: action receiver, control loop + + # FPS measurement + self.fps_tracker = FPSTracker(target_fps=self.config.fps) + + # Episode step counter (incremented in control_loop_action) + self._step_count = 0 + + self.logger.info("Robot connected and ready") + + # Use an event for thread-safe coordination + self.must_go = threading.Event() + self.must_go.set() # Initially set - observations qualify for direct processing + + @property + def running(self): + return not self.shutdown_event.is_set() + + def start(self): + """Start the robot client and connect to the policy server""" + try: + # client-server handshake + start_time = time.perf_counter() + self.stub.Ready(services_pb2.Empty()) + end_time = time.perf_counter() + self.logger.debug(f"Connected to policy server in {end_time - start_time:.4f}s") + + # send policy instructions + policy_config_bytes = pickle.dumps(self.policy_config) + policy_setup = services_pb2.PolicySetup(data=policy_config_bytes) + + self.logger.info("Sending policy instructions to policy server") + self.logger.debug( + f"Policy type: {self.policy_config.policy_type} | " + f"Pretrained name or path: {self.policy_config.pretrained_name_or_path} | " + f"Device: {self.policy_config.device}" + ) + + self.stub.SendPolicyInstructions(policy_setup) + + self.shutdown_event.clear() + + return True + + except grpc.RpcError as e: + self.logger.error(f"Failed to connect to policy server: {e}") + return False + + # ------------------------------------------------------------------ + # Home position / smooth reset (mirrors HTTP robot_server feature) + # ------------------------------------------------------------------ + + def move_to_home_slowly(self) -> None: + """Linearly interpolate from current joint positions to home over + reset_steps steps at reset_hz Hz. No-op if home_joints is empty.""" + if not self.config.home_joints: + return + + joint_names = list(self.robot.action_features.keys()) + home_values = [float(v) for v in self.config.home_joints.split()] + if len(home_values) != len(joint_names): + raise ValueError( + f"home_joints has {len(home_values)} values but robot has " + f"{len(joint_names)} joints: {joint_names}" + ) + + raw_obs = self.robot.get_observation() + current = [float(raw_obs[k]) for k in joint_names] + + dt = 1.0 / self.config.reset_hz + self.logger.info( + f"Moving to home over {self.config.reset_steps} steps " + f"at {self.config.reset_hz} Hz " + f"({self.config.reset_steps / self.config.reset_hz:.1f}s)" + ) + for i in range(1, self.config.reset_steps + 1): + t = i / self.config.reset_steps + interp = {k: c + t * (g - c) for k, c, g in zip(joint_names, current, home_values)} + self.robot.send_action(interp) + time.sleep(dt) + self.logger.info("Home position reached.") + + # ------------------------------------------------------------------ + # Episode helpers + # ------------------------------------------------------------------ + + def _stop_episode(self) -> None: + """End the current episode loop without disconnecting the robot.""" + self.shutdown_event.set() + + def reset_for_new_episode(self) -> None: + """Reset all in-memory state so a new episode can start cleanly. + The robot stays connected; call move_to_home_slowly() separately.""" + self.shutdown_event.clear() + self.action_queue = Queue() + self.action_queue_size = [] + with self.latest_action_lock: + self.latest_action = -1 + self.action_chunk_size = -1 + self.must_go = threading.Event() + self.must_go.set() + self._step_count = 0 + self.start_barrier = threading.Barrier(2) + self.fps_tracker.reset() + + def stop(self): + """Stop the robot client""" + self.shutdown_event.set() + + if self.teleop: + self.teleop.disconnect() + self.logger.debug("Teleoperator disconnected") + + self.robot.disconnect() + self.logger.debug("Robot disconnected") + + self.channel.close() + self.logger.debug("Client stopped, channel closed") + + def send_observation( + self, + obs: TimedObservation, + ) -> bool: + """Send observation to the policy server. + Returns True if the observation was sent successfully, False otherwise.""" + if not self.running: + raise RuntimeError("Client not running. Run RobotClient.start() before sending observations.") + + if not isinstance(obs, TimedObservation): + raise ValueError("Input observation needs to be a TimedObservation!") + + start_time = time.perf_counter() + observation_bytes = pickle.dumps(obs) + serialize_time = time.perf_counter() - start_time + self.logger.debug(f"Observation serialization time: {serialize_time:.6f}s") + + try: + observation_iterator = send_bytes_in_chunks( + observation_bytes, + services_pb2.Observation, + log_prefix="[CLIENT] Observation", + silent=True, + ) + _ = self.stub.SendObservations(observation_iterator) + obs_timestep = obs.get_timestep() + self.logger.debug(f"Sent observation #{obs_timestep} | ") + + return True + + except grpc.RpcError as e: + self.logger.error(f"Error sending observation #{obs.get_timestep()}: {e}") + return False + + def plan_direct(self, obs: TimedObservation) -> list[TimedAction]: + """Send an observation and receive actions in one synchronous gRPC call (Plan RPC). + + Alternative to the SendObservations + GetActions two-step. Blocks until the + server returns the planned actions. Returns an empty list on failure. + """ + observation_bytes = pickle.dumps(obs) + observation_iterator = send_bytes_in_chunks( + observation_bytes, + services_pb2.Observation, + log_prefix="[CLIENT] Plan", + silent=True, + ) + try: + actions_chunk = self.stub.Plan(observation_iterator) + if len(actions_chunk.data) == 0: + return [] + timed_actions: list[TimedAction] = pickle.loads(actions_chunk.data) # nosec + return timed_actions + except grpc.RpcError as e: + self.logger.error(f"Plan RPC error: {e}") + return [] + + def _inspect_action_queue(self): + with self.action_queue_lock: + queue_size = self.action_queue.qsize() + timestamps = sorted([action.get_timestep() for action in self.action_queue.queue]) + self.logger.debug(f"Queue size: {queue_size}, Queue contents: {timestamps}") + return queue_size, timestamps + + def _aggregate_action_queues( + self, + incoming_actions: list[TimedAction], + aggregate_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + ): + """Finds the same timestep actions in the queue and aggregates them using the aggregate_fn""" + if aggregate_fn is None: + # default aggregate function: take the latest action + def aggregate_fn(x1, x2): + return x2 + + future_action_queue = Queue() + with self.action_queue_lock: + internal_queue = self.action_queue.queue + + current_action_queue = {action.get_timestep(): action.get_action() for action in internal_queue} + + for new_action in incoming_actions: + with self.latest_action_lock: + latest_action = self.latest_action + + # New action is older than the latest action in the queue, skip it + if new_action.get_timestep() <= latest_action: + continue + + # If the new action's timestep is not in the current action queue, add it directly + elif new_action.get_timestep() not in current_action_queue: + future_action_queue.put(new_action) + continue + + # If the new action's timestep is in the current action queue, aggregate it + # TODO: There is probably a way to do this with broadcasting of the two action tensors + future_action_queue.put( + TimedAction( + timestamp=new_action.get_timestamp(), + timestep=new_action.get_timestep(), + action=aggregate_fn( + current_action_queue[new_action.get_timestep()], new_action.get_action() + ), + ) + ) + + with self.action_queue_lock: + self.action_queue = future_action_queue + + def receive_actions(self, verbose: bool = False): + """Receive actions from the policy server""" + # Wait at barrier for synchronized start + self.start_barrier.wait() + self.logger.info("Action receiving thread starting") + + while self.running: + try: + # Use StreamActions to get a stream of actions from the server + actions_chunk = self.stub.GetActions(services_pb2.Empty()) + if len(actions_chunk.data) == 0: + continue # received `Empty` from server, wait for next call + + receive_time = time.time() + + # Deserialize bytes back into list[TimedAction] + deserialize_start = time.perf_counter() + timed_actions = pickle.loads(actions_chunk.data) # nosec + deserialize_time = time.perf_counter() - deserialize_start + + # Log device type of received actions + if len(timed_actions) > 0: + received_device = timed_actions[0].get_action().device.type + self.logger.debug(f"Received actions on device: {received_device}") + + # Move actions to client_device (e.g., for downstream planners that need GPU) + client_device = self.config.client_device + if client_device != "cpu": + for timed_action in timed_actions: + if timed_action.get_action().device.type != client_device: + timed_action.action = timed_action.get_action().to(client_device) + self.logger.debug(f"Converted actions to device: {client_device}") + else: + self.logger.debug(f"Actions kept on device: {client_device}") + + self.action_chunk_size = max(self.action_chunk_size, len(timed_actions)) + + # Calculate network latency if we have matching observations + if len(timed_actions) > 0 and verbose: + with self.latest_action_lock: + latest_action = self.latest_action + + self.logger.debug(f"Current latest action: {latest_action}") + + # Get queue state before changes + old_size, old_timesteps = self._inspect_action_queue() + if not old_timesteps: + old_timesteps = [latest_action] # queue was empty + + # Log incoming actions + incoming_timesteps = [a.get_timestep() for a in timed_actions] + + first_action_timestep = timed_actions[0].get_timestep() + server_to_client_latency = (receive_time - timed_actions[0].get_timestamp()) * 1000 + + self.logger.info( + f"Received action chunk for step #{first_action_timestep} | " + f"Latest action: #{latest_action} | " + f"Incoming actions: {incoming_timesteps[0]}:{incoming_timesteps[-1]} | " + f"Network latency (server->client): {server_to_client_latency:.2f}ms | " + f"Deserialization time: {deserialize_time * 1000:.2f}ms" + ) + + # Update action queue + start_time = time.perf_counter() + self._aggregate_action_queues(timed_actions, self.config.aggregate_fn) + queue_update_time = time.perf_counter() - start_time + + self.must_go.set() # after receiving actions, next empty queue triggers must-go processing! + + if verbose: + # Get queue state after changes + new_size, new_timesteps = self._inspect_action_queue() + + with self.latest_action_lock: + latest_action = self.latest_action + + self.logger.info( + f"Latest action: {latest_action} | " + f"Old action steps: {old_timesteps[0]}:{old_timesteps[-1]} | " + f"Incoming action steps: {incoming_timesteps[0]}:{incoming_timesteps[-1]} | " + f"Updated action steps: {new_timesteps[0]}:{new_timesteps[-1]}" + ) + self.logger.debug( + f"Queue update complete ({queue_update_time:.6f}s) | " + f"Before: {old_size} items | " + f"After: {new_size} items | " + ) + + except grpc.RpcError as e: + self.logger.error(f"Error receiving actions: {e}") + + def actions_available(self): + """Check if there are actions available in the queue""" + with self.action_queue_lock: + return not self.action_queue.empty() + + def _action_tensor_to_action_dict(self, action_tensor: torch.Tensor) -> dict[str, float]: + action = {key: action_tensor[i].item() for i, key in enumerate(self.robot.action_features)} + return action + + def control_loop_action(self, verbose: bool = False) -> dict[str, Any]: + """Reading and performing actions in local queue""" + + # Lock only for queue operations + get_start = time.perf_counter() + with self.action_queue_lock: + self.action_queue_size.append(self.action_queue.qsize()) + # Get action from queue + timed_action = self.action_queue.get_nowait() + get_end = time.perf_counter() - get_start + + _performed_action = self.robot.send_action( + self._action_tensor_to_action_dict(timed_action.get_action()) + ) + with self.latest_action_lock: + self.latest_action = timed_action.get_timestep() + + # Episode step counter โ€” stop after max_episode_steps if set + self._step_count += 1 + if self.config.max_episode_steps and self._step_count >= self.config.max_episode_steps: + self.logger.info(f"Episode done: reached {self._step_count} steps.") + self._stop_episode() + + if verbose: + with self.action_queue_lock: + current_queue_size = self.action_queue.qsize() + + self.logger.debug( + f"Ts={timed_action.get_timestamp()} | " + f"Action #{timed_action.get_timestep()} performed | " + f"Queue size: {current_queue_size}" + ) + + self.logger.debug( + f"Popping action from queue to perform took {get_end:.6f}s | Queue size: {current_queue_size}" + ) + + return _performed_action + + def _ready_to_send_observation(self): + """Flags when the client is ready to send an observation""" + with self.action_queue_lock: + return self.action_queue.qsize() / self.action_chunk_size <= self._chunk_size_threshold + + def control_loop_observation(self, task: str, verbose: bool = False) -> RawObservation: + if not self.running: + return + try: + # Get serialized observation bytes from the function + start_time = time.perf_counter() + + raw_observation: RawObservation = self.robot.get_observation() + raw_observation["task"] = task + + with self.latest_action_lock: + latest_action = self.latest_action + + observation = TimedObservation( + timestamp=time.time(), # need time.time() to compare timestamps across client and server + observation=raw_observation, + timestep=max(latest_action + 1, 0), + ) + + obs_capture_time = time.perf_counter() - start_time + + # If there are no actions left in the queue, the observation must go through processing! + with self.action_queue_lock: + observation.must_go = self.must_go.is_set() and self.action_queue.empty() + current_queue_size = self.action_queue.qsize() + + _ = self.send_observation(observation) + + self.logger.debug(f"QUEUE SIZE: {current_queue_size} (Must go: {observation.must_go})") + if observation.must_go: + # must-go event will be set again after receiving actions + self.must_go.clear() + + if verbose: + # Calculate comprehensive FPS metrics + fps_metrics = self.fps_tracker.calculate_fps_metrics(observation.get_timestamp()) + + self.logger.info( + f"Obs #{observation.get_timestep()} | " + f"Avg FPS: {fps_metrics['avg_fps']:.2f} | " + f"Target: {fps_metrics['target_fps']:.2f}" + ) + + self.logger.debug( + f"Ts={observation.get_timestamp():.6f} | Capturing observation took {obs_capture_time:.6f}s" + ) + + return raw_observation + + except RuntimeError as e: + msg = str(e) + if "read thread is not running" in msg or "camera" in msg.lower(): + self.logger.warning(f"Camera error โ€” attempting reconnect: {e}") + try: + self.robot.disconnect() + time.sleep(1.0) + self.robot.connect() + self.logger.info("Robot reconnected successfully.") + except Exception as reconnect_err: + self.logger.error(f"Reconnect failed โ€” stopping episode: {reconnect_err}") + self._stop_episode() + else: + self.logger.error(f"Error in observation sender: {e}") + except Exception as e: + self.logger.error(f"Error in observation sender: {e}") + + def _teleop_control_loop(self, task: str, verbose: bool = False) -> None: + """Control loop for teleoperation mode. + + The leader arm drives the follower. Every step the observation + (including the actual teleop action under key 'teleop_action') is + streamed to the gRPC server so planner_server.py can record it. + No start_barrier is used โ€” there is no action-receiver thread. + """ + self.logger.info("Teleoperation control loop starting") + + while self.running: + loop_start = time.perf_counter() + + # 1. Read follower observation + try: + raw_obs: RawObservation = self.robot.get_observation() + except (ConnectionError, TimeoutError) as e: + self.logger.warning(f"Observation read failed, skipping: {e}") + time.sleep(self.config.environment_dt) + continue + + # 2. Get and apply leader action + teleop_action = self.teleop.get_action() + self.robot.send_action(teleop_action) + + # 3. Attach task + teleop action so the server can record them + raw_obs["task"] = task + raw_obs["teleop_action"] = list(teleop_action.values()) + + with self.latest_action_lock: + latest = self.latest_action + + obs = TimedObservation( + timestamp=time.time(), + observation=raw_obs, + timestep=max(latest + 1, 0), + must_go=True, # always send โ€” no action queue to pace against + ) + self.send_observation(obs) + + with self.latest_action_lock: + self.latest_action += 1 + + # 4. Episode step counter + self._step_count += 1 + if self.config.max_episode_steps and self._step_count >= self.config.max_episode_steps: + self.logger.info(f"Teleop episode done: reached {self._step_count} steps.") + self._stop_episode() + + if verbose: + fps_metrics = self.fps_tracker.calculate_fps_metrics(obs.get_timestamp()) + self.logger.info( + f"Teleop step #{self._step_count} | " + f"Avg FPS: {fps_metrics['avg_fps']:.2f}" + ) + + time.sleep(max(0.0, self.config.environment_dt - (time.perf_counter() - loop_start))) + + def control_loop(self, task: str, verbose: bool = False) -> tuple[Observation, Action]: + """Combined function for executing actions and streaming observations""" + # Wait at barrier for synchronized start + self.start_barrier.wait() + self.logger.info("Control loop thread starting") + + _performed_action = None + _captured_observation = None + + while self.running: + control_loop_start = time.perf_counter() + """Control loop: (1) Performing actions, when available""" + if self.actions_available(): + _performed_action = self.control_loop_action(verbose) + + """Control loop: (2) Streaming observations to the remote policy server""" + if self._ready_to_send_observation(): + _captured_observation = self.control_loop_observation(task, verbose) + + self.logger.debug(f"Control loop (ms): {(time.perf_counter() - control_loop_start) * 1000:.2f}") + # Dynamically adjust sleep time to maintain the desired control frequency + time.sleep(max(0, self.config.environment_dt - (time.perf_counter() - control_loop_start))) + + return _captured_observation, _performed_action + + +def _run_one_episode(client: "RobotClient") -> None: + """Run one episode โ€” teleop or policy mode.""" + if client.teleop: + # Teleop: single thread, no action receiver needed + client._teleop_control_loop(task=client.config.task) + else: + # Policy: action receiver + control loop threads + action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True) + action_receiver_thread.start() + client.control_loop(task=client.config.task) + action_receiver_thread.join(timeout=2.0) + + +@draccus.wrap() +def async_client(cfg: RobotClientConfig): + logging.info(pformat(asdict(cfg))) + + client = RobotClient(cfg) + + if not client.start(): + return + + episode_mode = cfg.max_episode_steps is not None or cfg.num_episodes is not None + num_episodes = cfg.num_episodes # None = run forever + + if not episode_mode: + # โ”€โ”€ Original continuous behaviour (unchanged) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + client.logger.info("Running in continuous mode (no episode limit).") + if client.teleop: + try: + client._teleop_control_loop(task=cfg.task) + finally: + client.stop() + client.logger.info("Client stopped") + else: + action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True) + action_receiver_thread.start() + try: + client.control_loop(task=cfg.task) + finally: + client.stop() + action_receiver_thread.join() + if cfg.debug_visualize_queue_size: + visualize_action_queue_size(client.action_queue_size) + client.logger.info("Client stopped") + + else: + # โ”€โ”€ Episode mode (new) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + ep = 0 + try: + while num_episodes is None or ep < num_episodes: + ep += 1 + client.logger.info(f"โ”€โ”€ Episode {ep} โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€") + if ep == 1: + client.move_to_home_slowly() + _run_one_episode(client) + if num_episodes is None or ep < num_episodes: + client.reset_for_new_episode() + # Block until the server finishes training on the completed episode. + client.logger.info("Waiting for server to complete training+reset...") + client.stub.Ready(services_pb2.Empty()) + client.move_to_home_slowly() + # Pause after arm is home โ€” robot holds still, no images sent. + if cfg.reset_pause_seconds > 0: + client.logger.info(f"Scene reset pause: {cfg.reset_pause_seconds:.0f}s ...") + time.sleep(cfg.reset_pause_seconds) + finally: + client.stop() + if cfg.debug_visualize_queue_size: + visualize_action_queue_size(client.action_queue_size) + client.logger.info(f"Done โ€” {ep} episode(s) completed.") + + +if __name__ == "__main__": + register_third_party_plugins() + async_client() # run the client diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/robot_server.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/robot_server.py new file mode 100644 index 0000000..db70cad --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/robot_server.py @@ -0,0 +1,216 @@ +""" +Minimal FastAPI robot server (port 7070) for the LeRobot side of the RL pipeline. + +The policy server (mpail2 env) calls these endpoints to drive the robot: + + GET /state โ†’ {"joints": [float, ...], "joint_names": [...]} + POST /step {"joints": [float, ...]} โ†’ execute joint targets (degrees) + POST /reset โ†’ move slowly to HOME_POSITION_DEG and return state + +Joint order is fixed by robot.action_features at startup and printed to the log. +The policy server must use the same order for every /step call. + +Usage: + python -m lerobot.async_inference.robot_server \\ + --robot.type=so100_follower \\ + --robot.port=/dev/ttyACM1 \\ + --robot.id=Kid_right \\ + --host=0.0.0.0 \\ + --port=7070 \\ + --home_joints="16.44 -104.44 95.91 72.75 -0.04 1.96" \\ + --reset_steps=80 \\ + --reset_hz=50 +""" + +import logging +import time +import threading +from contextlib import asynccontextmanager +from dataclasses import dataclass + +import draccus +import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel + +from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.robots import ( # noqa: F401 + RobotConfig, + bi_so_follower, + koch_follower, + make_robot_from_config, + omx_follower, + so_follower, +) +from lerobot.utils.import_utils import register_third_party_plugins + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger("robot_server") + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +@dataclass +class RobotServerConfig: + robot: RobotConfig + host: str = "0.0.0.0" + port: int = 7070 + # Space-separated home joint positions in degrees (order matches joint_names at startup). + home_joints: str = "16.44 -104.44 95.91 72.75 -0.04 1.96" + # Number of interpolation steps when moving to home (more = slower/smoother). + reset_steps: int = 80 + # Frequency (Hz) of each interpolation step during reset. + reset_hz: float = 50.0 + + +# --------------------------------------------------------------------------- +# Pydantic schemas +# --------------------------------------------------------------------------- + +class StepRequest(BaseModel): + joints: list[float] + +class StateResponse(BaseModel): + joints: list[float] + joint_names: list[str] + + +# --------------------------------------------------------------------------- +# Global robot state (set during lifespan startup) +# --------------------------------------------------------------------------- + +_robot = None +_joint_names: list[str] = [] # internal keys: "shoulder_pan.pos", ... +_joint_names_bare: list[str] = [] # bare names for API: "shoulder_pan", ... +_home_joints: list[float] = [] +_reset_steps: int = 80 +_reset_hz: float = 50.0 +_lock = threading.Lock() # serialize robot bus access + + +def _obs_to_joints(raw_obs: dict) -> list[float]: + return [float(raw_obs[k]) for k in _joint_names] + + +def _joints_to_action(joints: list[float]) -> dict: + return {k: v for k, v in zip(_joint_names, joints)} + + +def _move_to_slowly(target: list[float], steps: int, hz: float) -> list[float]: + """Interpolate from current position to target over `steps` steps at `hz` Hz.""" + dt = 1.0 / hz + raw_obs = _robot.get_observation() + current = _obs_to_joints(raw_obs) + + for i in range(1, steps + 1): + t = i / steps + interp = [c + t * (g - c) for c, g in zip(current, target)] + _robot.send_action(_joints_to_action(interp)) + time.sleep(dt) + + raw_obs = _robot.get_observation() + return _obs_to_joints(raw_obs) + + +# --------------------------------------------------------------------------- +# FastAPI app +# --------------------------------------------------------------------------- + +@asynccontextmanager +async def lifespan(app: FastAPI): + yield + if _robot is not None: + _robot.disconnect() + logger.info("Robot disconnected.") + + +app = FastAPI(lifespan=lifespan) + + +@app.get("/state", response_model=StateResponse) +def get_state(): + """Return current joint positions in degrees.""" + with _lock: + try: + raw_obs = _robot.get_observation() + except (ConnectionError, TimeoutError) as e: + raise HTTPException(status_code=503, detail=f"Robot read failed: {e}") + + joints = _obs_to_joints(raw_obs) + return StateResponse(joints=joints, joint_names=_joint_names_bare) + + +@app.post("/step", response_model=StateResponse) +def post_step(req: StepRequest): + """Execute joint targets (degrees) and return the resulting state.""" + if len(req.joints) != len(_joint_names): + raise HTTPException( + status_code=422, + detail=f"Expected {len(_joint_names)} joints, got {len(req.joints)}. " + f"Joint order: {_joint_names_bare}", + ) + + with _lock: + try: + _robot.send_action(_joints_to_action(req.joints)) + raw_obs = _robot.get_observation() + except (ConnectionError, TimeoutError) as e: + raise HTTPException(status_code=503, detail=f"Robot error: {e}") + + joints = _obs_to_joints(raw_obs) + return StateResponse(joints=joints, joint_names=_joint_names_bare) + + +@app.post("/reset", response_model=StateResponse) +def post_reset(): + """Move smoothly to home position and return resulting state.""" + with _lock: + try: + joints = _move_to_slowly(_home_joints, _reset_steps, _reset_hz) + except (ConnectionError, TimeoutError) as e: + raise HTTPException(status_code=503, detail=f"Robot error: {e}") + + logger.info(f"Reset to home: {dict(zip(_joint_names_bare, joints))}") + return StateResponse(joints=joints, joint_names=_joint_names_bare) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +@draccus.wrap() +def main(cfg: RobotServerConfig): + global _robot, _joint_names, _joint_names_bare, _home_joints, _reset_steps, _reset_hz + + register_third_party_plugins() + + _robot = make_robot_from_config(cfg.robot) + _robot.connect() + + _joint_names = list(_robot.action_features.keys()) + _joint_names_bare = [k.removesuffix(".pos") for k in _joint_names] + + home_values = [float(v) for v in cfg.home_joints.split()] + if len(home_values) != len(_joint_names): + raise ValueError( + f"--home_joints has {len(home_values)} values but robot has {len(_joint_names)} joints: " + f"{_joint_names_bare}" + ) + _home_joints = home_values + _reset_steps = cfg.reset_steps + _reset_hz = cfg.reset_hz + + logger.info(f"Joint order: {list(enumerate(_joint_names_bare))}") + logger.info(f"Home position: {dict(zip(_joint_names_bare, _home_joints))}") + logger.info(f"Reset motion: {cfg.reset_steps} steps at {cfg.reset_hz} Hz " + f"({cfg.reset_steps / cfg.reset_hz:.1f}s total)") + logger.info(f"Robot server starting on http://{cfg.host}:{cfg.port}") + + uvicorn.run(app, host=cfg.host, port=cfg.port, log_level="warning") + + +if __name__ == "__main__": + main() diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/teleop_with_planner.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/teleop_with_planner.py new file mode 100644 index 0000000..217cf49 --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/teleop_with_planner.py @@ -0,0 +1,176 @@ +""" +Teleoperate a robot while streaming observations to an external HTTP planner. + +The human stays in control (leader arm drives the follower). Observations are +forwarded to the planner at every step so the planner can observe what is +happening (e.g. for logging, imitation learning, or advisory output). +The planner's returned actions are logged but NOT applied to the robot. + +Usage (single arm): + python -m lerobot.async_inference.teleop_with_planner \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyUSB0 \ + --robot.id=my_robot \ + --teleop.type=so100_leader \ + --teleop.port=/dev/ttyUSB1 \ + --teleop.id=my_leader \ + --planner_url=http://localhost:9090/plan \ + --task="pick up the cup" \ + --fps=30 + +Notes: + - No bridge server (external_planner_server) is needed. + - Observations are POST-ed as JSON identical to the format external_planner_server uses. + - HTTP calls are fire-and-forget (non-blocking) so they never stall the control loop. +""" + +import json +import logging +import threading +import time +from dataclasses import dataclass + +import draccus +import numpy as np +import requests + +from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.robots import ( # noqa: F401 + RobotConfig, + bi_so_follower, + koch_follower, + make_robot_from_config, + omx_follower, + so_follower, +) +from lerobot.teleoperators import ( # noqa: F401 + TeleoperatorConfig, + bi_so_leader, + koch_leader, + make_teleoperator_from_config, + omx_leader, + so_leader, +) +from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.robot_utils import precise_sleep + +logger = logging.getLogger(__name__) + + +@dataclass +class TeleopWithPlannerConfig: + robot: RobotConfig + teleop: TeleoperatorConfig + planner_url: str = "http://localhost:9090/plan" + task: str = "teleoperation" + fps: int = 30 + # Seconds before giving up on a planner HTTP response + http_timeout: float = 2.0 + # Whether to include camera images in the JSON payload (large!) + include_images: bool = False + + +def _build_payload( + raw_obs: dict, + timestep: int, + task: str, + include_images: bool, + teleop_action: dict | None = None, +) -> dict: + observation_payload: dict = {"task": task} + for key, value in raw_obs.items(): + if isinstance(value, np.ndarray): + if not include_images and value.ndim == 3: + continue + observation_payload[key] = value.tolist() + elif hasattr(value, "tolist"): + observation_payload[key] = value.tolist() + else: + observation_payload[key] = value + + payload: dict = { + "timestep": timestep, + "timestamp": time.time(), + "must_go": False, + "observation": observation_payload, + } + # Include actual teleop action so planner_server records what the human did, + # not the planner's own prediction. + if teleop_action is not None: + payload["teleop_action"] = [float(v) for v in teleop_action.values()] + return payload + + +def _post_to_planner(url: str, payload: dict, timeout: float) -> None: + """Called in a background thread โ€” never blocks the control loop.""" + try: + resp = requests.post( + url, + data=json.dumps(payload), + headers={"Content-Type": "application/json"}, + timeout=timeout, + ) + resp.raise_for_status() + data = resp.json() + actions = data.get("actions", []) + logger.debug(f"Planner returned {len(actions)} action(s) for step #{payload['timestep']}") + except requests.RequestException as e: + logger.warning(f"Planner HTTP error at step #{payload['timestep']}: {e}") + + +@draccus.wrap() +def main(cfg: TeleopWithPlannerConfig): + logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + + robot = make_robot_from_config(cfg.robot) + teleop = make_teleoperator_from_config(cfg.teleop) + + robot.connect() + teleop.connect() + + logger.info(f"Teleoperating robot โ€” streaming observations to {cfg.planner_url}") + logger.info("Press Ctrl+C to stop.") + + environment_dt = 1.0 / cfg.fps + timestep = 0 + + try: + while True: + loop_start = time.perf_counter() + + # 1. Capture observation from follower arm + try: + raw_obs = robot.get_observation() + raw_obs["task"] = cfg.task + except (ConnectionError, TimeoutError) as e: + logger.warning(f"Observation read failed at step #{timestep}, skipping: {e}") + timestep += 1 + precise_sleep(max(0.0, environment_dt - (time.perf_counter() - loop_start))) + continue + + # 2. Get teleop action from leader arm and apply it to the follower + action = teleop.get_action() + robot.send_action(action) + + # 3. Fire-and-forget POST to the planner (non-blocking) + payload = _build_payload(raw_obs, timestep, cfg.task, cfg.include_images, action) + threading.Thread( + target=_post_to_planner, + args=(cfg.planner_url, payload, cfg.http_timeout), + daemon=True, + ).start() + + timestep += 1 + precise_sleep(max(0.0, environment_dt - (time.perf_counter() - loop_start))) + + except KeyboardInterrupt: + logger.info("Stopped.") + finally: + robot.disconnect() + teleop.disconnect() + + +if __name__ == "__main__": + register_third_party_plugins() + main() diff --git a/mpail2/envs/real/so101/lerobot_patch/async_inference/test_communication.py b/mpail2/envs/real/so101/lerobot_patch/async_inference/test_communication.py new file mode 100644 index 0000000..f31209d --- /dev/null +++ b/mpail2/envs/real/so101/lerobot_patch/async_inference/test_communication.py @@ -0,0 +1,166 @@ +""" +Minimal communication test โ€” no real robot needed. + +Run in THREE separate terminals: + + Terminal 1 (mock external planner, HTTP): + python -m lerobot.async_inference.test_communication planner + + Terminal 2 (bridge server): + python -m lerobot.async_inference.test_communication bridge + + Terminal 3 (fake robot that sends one observation and prints the action it gets back): + python -m lerobot.async_inference.test_communication robot +""" + +import json +import pickle +import sys +import time +from http.server import BaseHTTPRequestHandler, HTTPServer + +import grpc +import numpy as np + +from lerobot.transport import services_pb2, services_pb2_grpc +from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks + +from .helpers import TimedObservation + + +# --------------------------------------------------------------------------- +# Terminal 1: mock external planner (HTTP server) +# --------------------------------------------------------------------------- + +class MockPlannerHandler(BaseHTTPRequestHandler): + """Receives an observation JSON, prints it, returns dummy actions.""" + + def do_POST(self): # noqa: N802 + length = int(self.headers["Content-Length"]) + body = self.rfile.read(length) + obs = json.loads(body) + + print("\n[PLANNER] Got observation:") + print(f" timestep : {obs['timestep']}") + print(f" must_go : {obs['must_go']}") + state = obs["observation"].get("observation.state", []) + print(f" state : {state}") + task = obs["observation"].get("task", "") + print(f" task : {task}") + + # Return 3 dummy actions (same dimensionality as the fake state) + action_dim = len(state) if state else 6 + actions = [[float(i) * 0.01] * action_dim for i in range(3)] + response = json.dumps({"actions": actions}).encode() + + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + print("[PLANNER] Sent back 3 dummy actions.\n") + + def log_message(self, *args): # suppress default access log + pass + + +def run_planner(host="0.0.0.0", port=9090): + print(f"[PLANNER] Mock external planner listening on {host}:{port}") + HTTPServer((host, port), MockPlannerHandler).serve_forever() + + +# --------------------------------------------------------------------------- +# Terminal 2: bridge server (ExternalPlannerServer) +# --------------------------------------------------------------------------- + +def run_bridge(host="0.0.0.0", port=8080, planner_url="http://127.0.0.1:9090/plan"): + from lerobot.async_inference.external_planner_server import serve + print(f"[BRIDGE] Starting ExternalPlannerServer on {host}:{port}") + print(f"[BRIDGE] Will forward observations to {planner_url}") + serve(host=host, port=port, planner_url=planner_url) + + +# --------------------------------------------------------------------------- +# Terminal 3: fake robot client (sends one observation, prints action) +# --------------------------------------------------------------------------- + +def run_robot(server_address="127.0.0.1:8080"): + print(f"[ROBOT] Connecting to bridge server at {server_address}") + channel = grpc.insecure_channel(server_address, grpc_channel_options()) + stub = services_pb2_grpc.AsyncInferenceStub(channel) + + # 1. Handshake + stub.Ready(services_pb2.Empty()) + print("[ROBOT] Handshake OK") + + # 2. Send fake policy instructions (bridge logs them, doesn't use them) + from .helpers import RemotePolicyConfig + fake_config = RemotePolicyConfig( + policy_type="act", + pretrained_name_or_path="test/model", + lerobot_features={}, + actions_per_chunk=3, + device="cpu", + ) + stub.SendPolicyInstructions(services_pb2.PolicySetup(data=pickle.dumps(fake_config))) + print("[ROBOT] Policy instructions sent") + + # 3. Build a fake observation (6-DOF joint state) + fake_raw_obs = { + "observation.state": np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], dtype=np.float32), + "task": "test communication", + } + timed_obs = TimedObservation( + timestamp=time.time(), + timestep=0, + observation=fake_raw_obs, + must_go=True, + ) + + obs_bytes = pickle.dumps(timed_obs) + obs_iterator = send_bytes_in_chunks(obs_bytes, services_pb2.Observation, silent=True) + stub.SendObservations(obs_iterator) + print("[ROBOT] Observation sent") + + # 4. Request actions back + print("[ROBOT] Waiting for actions...") + actions_msg = stub.GetActions(services_pb2.Empty()) + + if not actions_msg.data: + print("[ROBOT] Received empty response (timeout or error on bridge/planner side)") + else: + timed_actions = pickle.loads(actions_msg.data) # nosec + print(f"[ROBOT] Received {len(timed_actions)} actions:") + for a in timed_actions: + print(f" timestep={a.get_timestep()} action={a.get_action().tolist()}") + + channel.close() + print("[ROBOT] Done.") + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +USAGE = """ +Usage: + python -m lerobot.async_inference.test_communication planner # Terminal 1 + python -m lerobot.async_inference.test_communication bridge # Terminal 2 + python -m lerobot.async_inference.test_communication robot # Terminal 3 +""" + +if __name__ == "__main__": + if len(sys.argv) < 2: + print(USAGE) + sys.exit(1) + + role = sys.argv[1] + if role == "planner": + run_planner() + elif role == "bridge": + run_bridge() + elif role == "robot": + run_robot() + else: + print(f"Unknown role '{role}'.{USAGE}") + sys.exit(1) diff --git a/mpail2/envs/real/so101/network/__init__.py b/mpail2/envs/real/so101/network/__init__.py new file mode 100644 index 0000000..4e49840 --- /dev/null +++ b/mpail2/envs/real/so101/network/__init__.py @@ -0,0 +1 @@ +"""gRPC server owning the physical SO-101 arm + cameras.""" diff --git a/mpail2/envs/real/so101/network/server.py b/mpail2/envs/real/so101/network/server.py new file mode 100644 index 0000000..9d35e3d --- /dev/null +++ b/mpail2/envs/real/so101/network/server.py @@ -0,0 +1,355 @@ +""" +so101_robot_server.py โ€” gRPC control server exposing the real SO-101 arm + cameras +to mpail2/envs/real/so101/so101_env.py. + +Runs directly against lerobot's SOFollower driver โ€” no LeRobot async_inference +layer involved, no separate client process required on this side. Mirrors +mpail2/envs/real/franka/network/server.py's role (thin RPC wrapper around the +hardware driver), but gRPC instead of ZMQ, and combines action+observation into +one round-trip per step instead of separate reads. + +Run in the `lerobot` conda env (needs `mpail2` installed there too, e.g. `pip install -e ".[so101]"`): + + conda activate lerobot + python -m mpail2.envs.real.so101.network.server \ + --robot_port /dev/ttyACM0 \ + --robot_id Kid \ + --cam_index /dev/video0 \ + --cam2_serial 317422074482 \ + --grpc_port 7070 + +Then point the mpail2-side env at it (mock=False): + + from mpail2.envs.real.so101 import SO101RealEnvArgs, make_so101_env + env = make_so101_env(SO101RealEnvArgs(host="", port=7070, mock=False)) +""" + +import argparse +import logging +import time +from concurrent import futures + +import grpc +import numpy as np + +from mpail2.envs.real.so101.transport import so101_robot_pb2, so101_robot_pb2_grpc + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +logger = logging.getLogger("so101_robot_server") + +# Must match mpail2/envs/real/so101/robot_limits.py โ€” JOINT_NAMES, HOME_POSITION_DEG, +# CAM_KEY/CAM2_KEY, CAM_H/CAM_W/CAM2_H/CAM2_W. Kept as local constants here since this +# script runs in the separate `lerobot` conda env, which doesn't have mpail2 installed. +JOINT_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"] +HOME_POSITION_DEG = [-21.19, -5.41, 1.58, 99.47, -14.20, 28.66] +CAM_KEY, CAM2_KEY = "cam", "cam2" +# 64x48 preserves the native 640x480 (4:3) aspect ratio exactly (10x isotropic downsample, +# no stretch) โ€” unlike 84x84, which squashes x/y by different factors (84/640 vs 84/480). +CAM_W, CAM_H = 64, 48 +CAM2_W, CAM2_H = 64, 48 + + +def _joints_dict_to_array(obs: dict) -> np.ndarray: + return np.array([obs[f"{name}.pos"] for name in JOINT_NAMES], dtype=np.float32) + + +def _to_bytes(img: np.ndarray, h: int, w: int) -> bytes: + """uint8 (H0, W0, 3) camera frame -> raw bytes at (h, w, 3), resizing if needed.""" + img = np.asarray(img) + if img.shape[0] != h or img.shape[1] != w: + import cv2 + img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA) + return np.ascontiguousarray(img, dtype=np.uint8).tobytes() + + +class SO101RobotServicer(so101_robot_pb2_grpc.SO101RobotServicer): + def __init__( + self, robot, settle_tolerance_deg: float = 0.5, settle_timeout_s: float = 0.5, + gripper_settle_tolerance_deg: float = 3.0, + reset_duration_s: float = 4.0, reset_hz: float = 20.0, + ): + self.robot = robot + self.settle_tolerance_deg = settle_tolerance_deg + self.settle_timeout_s = settle_timeout_s + self.reset_duration_s = reset_duration_s + self.reset_hz = reset_hz + # Gripper is a much slower/weaker actuator than the 5 arm joints (and often fights real + # resistance โ€” squeezing an object, brushing the table) โ€” holding it to the same 1ยฐ arm + # tolerance means it never settles, so every Step() burns the full timeout every time, + # silently undoing whatever control rate was requested. It still gets commanded and its + # readback is still returned; it's just not what the wait blocks on. + self.per_joint_tolerance_deg = {name: settle_tolerance_deg for name in JOINT_NAMES} + self.per_joint_tolerance_deg["gripper"] = gripper_settle_tolerance_deg + + def _state_from_obs(self, obs: dict) -> "so101_robot_pb2.RobotState": + joints = _joints_dict_to_array(obs) + cam = _to_bytes(obs[CAM_KEY], CAM_H, CAM_W) if CAM_KEY in obs else b"" + cam2 = _to_bytes(obs[CAM2_KEY], CAM2_H, CAM2_W) if CAM2_KEY in obs else b"" + return so101_robot_pb2.RobotState(joints_deg=joints.tolist(), cam=cam, cam2=cam2) + + def _read_present_position(self, retries: int = 5, backoff_s: float = 0.005) -> dict: + """bus.sync_read('Present_Position') with retry-on-transient-error. + + The half-duplex serial bus occasionally returns a corrupted status packet under + load (raises ConnectionError) โ€” with the settle-wait polling this every 10ms, + that's a "when", not "if". Absorb a few transient glitches here rather than + letting one crash the whole Step()/Reset() RPC; if it's a real disconnect + (not transient), it'll still raise after exhausting retries. + """ + for attempt in range(retries): + try: + return self.robot.bus.sync_read("Present_Position") + except ConnectionError: + if attempt == retries - 1: + raise + time.sleep(backoff_s) + + def _get_observation_sync(self) -> dict: + """Like robot.get_observation(), but calls cam.read() (blocking โ€” waits for the + NEXT frame captured after this call starts) instead of cam.read_latest() + (non-blocking peek at whatever's already sitting in the background capture + thread's buffer, which can be up to ~1/fps stale). Guarantees the returned + frames were captured after the joint read below, at the cost of blocking for + up to ~1/fps per camera instead of returning instantly. Always used โ€” see + so101_env.py/grpc_policy_server.py, which both expect the policy to act on a + freshly-captured observation rather than a possibly-stale buffered one. + """ + obs = {f"{name}.pos": val for name, val in self._read_present_position().items()} + for cam_key, cam in self.robot.cameras.items(): + obs[cam_key] = cam.read() + return obs + + def _get_observation_safe(self, retries: int = 3, backoff_s: float = 0.005) -> dict: + """_get_observation_sync() with retry-on-transient-error.""" + for attempt in range(retries): + try: + return self._get_observation_sync() + except ConnectionError: + if attempt == retries - 1: + raise + time.sleep(backoff_s) + + def _wait_until_settled(self, goal_pos: dict) -> tuple[bool, dict]: + """Block until every joint's readback is within tolerance of its goal, or timeout. + + send_action()/sync_write() is fire-and-forget โ€” it returns as soon as the goal + register is written, not once the servo arrives. At high step rates the very + next Step() would otherwise overwrite the goal before heavier joints (shoulder/ + elbow/wrist, real inertia) finish traveling, so get_observation() reads a stale, + frozen-looking position. Poll Present_Position directly (cheaper than a full + get_observation(), which also grabs both camera frames) until settled. + + settle_timeout_s <= 0 means wait indefinitely โ€” no deadline is set at all. Only + safe if per_joint_tolerance_deg is loose enough for every joint to actually reach; + a joint with a persistent steady-state offset (e.g. gravity droop) that never + closes will otherwise block this call forever. + + Returns (settled, last_errors) โ€” last_errors is {joint: |present-goal|} from the + final poll, so callers can log how far off a timeout actually was (a fraction of + a degree off vs. stuck 20 degrees away are very different problems). + """ + deadline = time.time() + self.settle_timeout_s if self.settle_timeout_s > 0 else None + while True: + present = self._read_present_position() + errors = {name: abs(present[name] - goal_pos[name]) for name in goal_pos} + if all(err <= self.per_joint_tolerance_deg[name] for name, err in errors.items()): + return True, errors + if deadline is not None and time.time() >= deadline: + return False, errors + time.sleep(0.01) + + def Reset(self, request, context): + # Ease into home instead of one big instant jump: send a ramp of intermediate + # targets (linear interpolation, present -> home) at reset_hz, over reset_duration_s. + # A single send_action(HOME_POSITION_DEG) commands the servos at full speed/torque + # from wherever they currently are, which can be a large, abrupt, jerky move if the + # previous episode ended far from home. + present = self._read_present_position() + start = {name: present[name] for name in JOINT_NAMES} + goal = dict(zip(JOINT_NAMES, HOME_POSITION_DEG)) + + n_steps = max(1, int(self.reset_duration_s * self.reset_hz)) + dt = 1.0 / self.reset_hz + for i in range(1, n_steps + 1): + alpha = i / n_steps + interp = {name: start[name] + alpha * (goal[name] - start[name]) for name in JOINT_NAMES} + self.robot.send_action({f"{name}.pos": val for name, val in interp.items()}) + time.sleep(dt) + + self._wait_until_settled(goal) + obs = self._get_observation_safe() + logger.info(f"Reset -> home position (ramped over {self.reset_duration_s:.1f}s)") + return self._state_from_obs(obs) + + def Step(self, request, context): + joints = list(request.joints_deg) + if len(joints) != len(JOINT_NAMES): + context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + f"Expected {len(JOINT_NAMES)} joint targets, got {len(joints)}", + ) + action = dict(zip([f"{n}.pos" for n in JOINT_NAMES], joints)) + sent = self.robot.send_action(action) + goal_pos = {k.removesuffix(".pos"): v for k, v in sent.items()} + settled, errors = self._wait_until_settled(goal_pos) + if not settled: + over_tol = { + name: err for name, err in errors.items() + if err > self.per_joint_tolerance_deg[name] + } + worst = sorted(over_tol.items(), key=lambda kv: -kv[1]) + worst_str = ", ".join(f"{name}={err:.2f}ยฐ" for name, err in worst) + logger.warning( + f"Step: arm did not settle within {self.settle_timeout_s:.2f}s โ€” " + f"reading present position anyway. Over tolerance: {worst_str}" + ) + obs = self._get_observation_safe() + return self._state_from_obs(obs) + + +def serve(): + parser = argparse.ArgumentParser(description="gRPC control server for the real SO-101 arm") + parser.add_argument("--robot_port", default="/dev/ttyACM0", help="Serial port the arm is connected to") + parser.add_argument("--robot_id", default="Kid", help="Calibration id (matches the LeRobot calibration file)") + parser.add_argument("--cam_index", default="/dev/video0", help="Wrist cam (opencv) index or device path") + parser.add_argument("--cam2_serial", default=None, help="RealSense serial number for cam2 (omit to disable)") + parser.add_argument("--cam_fps", type=int, default=30, + help="Capture fps for both cameras (default: 30).") + parser.add_argument("--grpc_host", default="[::]") + parser.add_argument("--grpc_port", type=int, default=7070) + parser.add_argument( + "--max_relative_target", type=float, default=None, + help="Per-step delta clamp passed straight to lerobot's SOFollowerRobotConfig (default: no clamp).", + ) + parser.add_argument( + "--settle_tolerance_deg", type=float, default=2.0, + help="Step()/Reset() block until every arm joint (all but gripper) is within this many " + "degrees of its goal (or --settle_timeout_s elapses) before reading the observation back.", + ) + parser.add_argument( + "--gripper_settle_tolerance_deg", type=float, default=15.0, + help="Separate, looser settle tolerance for the gripper โ€” it's a slower/weaker actuator " + "than the arm joints (and often fights real resistance), so holding it to the arm's " + "tolerance means it never settles and every step burns the full timeout for nothing.", + ) + parser.add_argument( + "--settle_timeout_s", type=float, default=0.5, + help="Max time to wait for the arm to settle before giving up and reading present position " + "anyway. <= 0 disables the timeout entirely (wait indefinitely) โ€” only safe if " + "settle/gripper tolerances are loose enough for every joint to actually be reachable, " + "or a joint with a persistent steady-state offset will hang Step() forever.", + ) + parser.add_argument( + "--reset_duration_s", type=float, default=4.0, + help="Reset() ramps from the current position to home over this many seconds " + "(linear interpolation), instead of commanding home in one instant jump.", + ) + parser.add_argument( + "--reset_hz", type=float, default=20.0, + help="Update rate of the Reset() ramp's intermediate targets.", + ) + parser.add_argument( + "--p_coefficient", type=int, default=16, + help="Position-loop proportional gain applied to every joint after connect, overriding " + "LeRobot's own default of 16 ('to avoid shakiness', tuned for smooth human-paced " + "teleop). RL-driven targets change faster/less smoothly than a teleop operator's hand, " + "so the softer gain isn't enough torque to close small residual errors in time โ€” this " + "was the confirmed root cause of the persistent settle-timeout warnings. 32 restores " + "Feetech's own factory default. Set higher for tighter tracking at the risk of " + "oscillation/shakiness; set to 16 (or omit this override) to match stock teleop behavior.", + ) + parser.add_argument( + "--gripper_torque_limit", type=int, default=500, + help="Max_Torque_Limit for the gripper only (0-1000 = 0-100%%), overriding LeRobot's own " + "default of 500 (50%%, 'to avoid burnout'). 800 = 80%%, a meaningful increase while " + "leaving headroom below full torque given the gripper is the joint most prone to " + "sustained resistance (squeezing objects, hitting the table). Raise further only with " + "awareness of overheating risk under prolonged stall.", + ) + parser.add_argument( + "--i_coefficient", type=int, default=4, + help="Position-loop integral gain applied to every joint after connect, overriding " + "LeRobot's own default of 0 (pure P+D, no integral term at all). A P+D-only loop " + "cannot drive steady-state error to zero under any constant load (gravity on " + "shoulder_lift, static friction, gripper squeezing an object) โ€” it settles into a " + "stable equilibrium where the P term's output exactly balances the load and simply " + "stops there, which is why raising --p_coefficient alone and/or waiting longer " + "(--settle_timeout_s) doesn't help: the system already reached its (wrong) equilibrium " + "and isn't converging further at all. A small nonzero I term lets the controller " + "accumulate that persistent error over time and keep increasing output until it's " + "actually eliminated. Start small and increase cautiously โ€” too high an I gain causes " + "integral windup / overshoot / oscillation. 0 restores stock teleop behavior.", + ) + parser.add_argument( + "--goal_velocity", type=int, default=0, + help="Goal_Velocity register (address 46) written to every motor after connect โ€” a " + "hardware-level speed cap: the servo's own control loop will not move faster than " + "this toward whatever Goal_Position it's given, no matter how large the commanded " + "jump is. This is the same kind of protection Kinova/Franka get for free from their " + "native Cartesian controllers; these Feetech servos don't have it unless set " + "explicitly. 0 = unlimited (stock/factory behavior, no cap). Units are raw register " + "ticks, not degrees/sec directly โ€” tune empirically by watching actual motion " + "smoothness, there's no documented tick-to-deg/s conversion in this codebase yet.", + ) + args = parser.parse_args() + + from lerobot.cameras.opencv import OpenCVCameraConfig + from lerobot.robots.so_follower import SO101Follower, SOFollowerRobotConfig + + cameras = {CAM_KEY: OpenCVCameraConfig(index_or_path=args.cam_index, width=640, height=480, fps=args.cam_fps)} + if args.cam2_serial: + from lerobot.cameras.realsense import RealSenseCameraConfig + cameras[CAM2_KEY] = RealSenseCameraConfig( + serial_number_or_name=args.cam2_serial, width=640, height=480, fps=args.cam_fps + ) + + config = SOFollowerRobotConfig( + port=args.robot_port, + id=args.robot_id, + cameras=cameras, + max_relative_target=args.max_relative_target, + ) + robot = SO101Follower(config) + logger.info(f"Connecting to SO-101 on {args.robot_port} (id={args.robot_id})...") + robot.connect() + logger.info("Connected.") + + # robot.connect() -> configure() just wrote LeRobot's own teleop-tuned defaults + # (P_Coefficient=16, I_Coefficient=0, gripper Max_Torque_Limit=500) โ€” override for this + # RL-driven training/eval server, where per-step targets change faster/less smoothly than + # a teleop operator's hand. I_Coefficient=0 (pure P+D) is the confirmed root cause of the + # persistent settle-timeout warnings: with no integral term the loop settles into a stable + # equilibrium under any constant load (gravity, friction, gripper resistance) instead of + # ever fully closing the error โ€” no amount of extra P gain or settle_timeout_s fixes that, + # since the system isn't still converging, it's already stopped. + for motor in robot.bus.motors: + robot.bus.write("P_Coefficient", motor, args.p_coefficient) + robot.bus.write("I_Coefficient", motor, args.i_coefficient) + robot.bus.write("Goal_Velocity", motor, args.goal_velocity) + robot.bus.write("Max_Torque_Limit", "gripper", args.gripper_torque_limit) + logger.info( + f"Overrode servo tuning for training: P_Coefficient={args.p_coefficient}, " + f"I_Coefficient={args.i_coefficient}, Goal_Velocity={args.goal_velocity} (all joints), " + f"gripper Max_Torque_Limit={args.gripper_torque_limit}/1000" + ) + + server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + servicer = SO101RobotServicer( + robot, settle_tolerance_deg=args.settle_tolerance_deg, settle_timeout_s=args.settle_timeout_s, + gripper_settle_tolerance_deg=args.gripper_settle_tolerance_deg, + reset_duration_s=args.reset_duration_s, reset_hz=args.reset_hz, + ) + so101_robot_pb2_grpc.add_SO101RobotServicer_to_server(servicer, server) + server.add_insecure_port(f"{args.grpc_host}:{args.grpc_port}") + server.start() + logger.info(f"so101_robot_server listening on {args.grpc_host}:{args.grpc_port}") + try: + server.wait_for_termination() + except KeyboardInterrupt: + logger.info("Shutting down...") + robot.disconnect() + server.stop(grace=2.0) + + +if __name__ == "__main__": + serve() diff --git a/mpail2/envs/real/so101/robot_limits.py b/mpail2/envs/real/so101/robot_limits.py new file mode 100644 index 0000000..0ae3d75 --- /dev/null +++ b/mpail2/envs/real/so101/robot_limits.py @@ -0,0 +1,117 @@ +"""Physical constants and defaults for the SO-101 (SO-ARM101) 6-DOF robot arm. + +Action space: Cartesian end-effector (x, y, z) + wrist_roll + gripper โ€” 5-dim. +State space: 6 joint positions in degrees (unchanged). +IK converts the 5-dim policy action back to 6 joint targets before sending to the robot. +""" + +import numpy as np + +# โ”€โ”€โ”€ dimensions โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +STATE_DIM = 6 # 6 joint positions reported by LeRobot (raw, used for IK only) +ACTION_DIM = 5 # [x, y, z, wrist_roll, gripper] โ€” x,y,z is an ABSOLUTE target + # position (mapped from [-1,1] onto the EE_LOWER_M/EE_UPPER_M box), + # not a delta. wrist_roll and gripper are still delta/increment style. +EE_PROPRIO_DIM = 13 # proprioception: [ee_x, ee_y, ee_z, qx, qy, qz, qw, j0..j5] + +# Joint names in the order LeRobot reports / planner_server.py parses them +JOINT_NAMES = [ + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", + "gripper", +] + +# Joint position bounds (degrees). Tune to your calibration. +JOINT_LOWER_DEG = np.array([ + -51.16, # shoulder_pan + -26.02, # shoulder_lift โ€” home(-6.022) - 20 + -90.0, # elbow_flex โ€” demo reaches -82.9ยฐ; extended from -17.1 + -66.77, # wrist_flex + -60.0, # wrist_roll + 3.0, # gripper โ€” current demo.pt reaches 4.02ยฐ (1st pct 4.02ยฐ); lowered from 9.0, + # which was clipping ~5ยฐ of the demo's full-closed range every step. +], dtype=np.float32) + +JOINT_UPPER_DEG = np.array([ + 35.0, # shoulder_pan + 65.0, # shoulder_lift โ€” demo reaches 62.2ยฐ; extended from 13.98 + 40.0, # elbow_flex โ€” current demo.pt reaches 38.68ยฐ (99th pct 31.05ยฐ); raised from + # 27.0, which was clipping ~12ยฐ off the demo's most-bent elbow poses. + 104.0, # wrist_flex โ€” demo reaches 103.30ยฐ; raised from 102.5 + 60.0, # wrist_roll + 120.0, # gripper +], dtype=np.float32) + +# Home position (degrees) โ€” calibrated from physical robot. +# Order matches JOINT_NAMES: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper +HOME_POSITION_DEG = np.array([-21.19, -5.41, 1.58, 99.47, -14.20, 28.66], dtype=np.float32) + +# โ”€โ”€โ”€ Cartesian workspace โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# Home EE position computed from FK at HOME_POSITION_DEG (arm joints only). +HOME_EE_M = np.array([0.1861, 0.0592, 0.0878], dtype=np.float32) # metres + +# NOTE: so101_env.py's xyz action is now an ABSOLUTE position (mapped from [-1,1] +# onto the EE_LOWER_M/EE_UPPER_M box below), not a delta โ€” MAX_DELTA_M is unused +# there. Still used by grpc_policy_server.py's (separate, untouched) delta-style +# action mapping: action_norm โˆˆ [-1,1]^3 is a per-step displacement scaled by +# MAX_DELTA_M * speed_scale, action=0 โ†’ hold current position. +MAX_DELTA_M = np.float32(0.03) # metres per step at speed_scale=1.0 + +# Hard workspace bounds โ€” EE is clipped to this box (xyz is an absolute position target, +# see so101_env.py's _action_norm_to_joints). Tightened to the current demo.pt's actual +# reach (recomputed via FK over all 5970 transitions, 2026-08-10 recording): +# x: min=0.1619 1st pct=0.1666 99th pct=0.2618 max=0.2782 +# y: min=-0.0726 1st pct=-0.0657 99th pct=0.0586 max=0.0598 +# z: min=-0.0141 1st pct=-0.0098 99th pct=0.0905 max=0.0923 +# x upper raised 0.24 -> 0.28 (demo max 0.2782 + margin) โ€” the previous bound was +# clipping ~9.6% of this demo's transitions before the target EE position was ever +# reached, misaligning the demo's target with what the live env could actually hit. +EE_LOWER_M = np.array([0.14, -0.08, -0.02], dtype=np.float32) +EE_UPPER_M = np.array([0.28, 0.09, 0.10], dtype=np.float32) + +# Kept for reference (not used in action mapping anymore). +EE_HALF_RANGE_M = (EE_UPPER_M - EE_LOWER_M) / 2 + +# Per-joint speed limit in degrees (legacy, used by so101_env.py joint-space path). +MAX_DELTA_DEG = np.array([10, 10, 10, 10, 10, 30], dtype=np.float32) + +# Wrist roll action component (index 3): direct degree control. +HOME_WRIST_ROLL_DEG = np.float32(-14.20) # matches HOME_POSITION_DEG[4] +WRIST_ROLL_HALF_RANGE = np.float32(5.0) # max ยฐ/step at speed_scale=1.0 (delta action) + +# Gripper action component (index 4): direct degree control. +# Not read by the actual clip (so101_env.py clips against JOINT_LOWER_DEG[5]/ +# JOINT_UPPER_DEG[5] instead) โ€” kept in sync here so this isn't a stale/misleading +# duplicate for anything that does reference it. +GRIPPER_LOWER_DEG = np.float32(3.0) +GRIPPER_UPPER_DEG = np.float32(120.0) +HOME_GRIPPER_DEG = np.float32(28.66) # matches HOME_POSITION_DEG[5] +GRIPPER_HALF_RANGE = np.float32(30.0) # max ยฐ/step at speed_scale=1.0 (delta action) โ€” demo closes ~44ยฐ in 20 steps + +# โ”€โ”€โ”€ timing โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +MAX_EPISODE_STEPS = 200 # 300 client steps โ†’ 299 transitions, fills storage exactly +CONTROL_FREQUENCY_HZ = 10.0 # env step rate + +# โ”€โ”€โ”€ LeRobot robot-control HTTP server โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# The env talks to a lightweight HTTP server running on the LeRobot side. +# Implement the server (see so101_env.py docstring for the expected API). +DEFAULT_ROBOT_HOST = "127.0.0.1" +DEFAULT_ROBOT_PORT = 7070 # intentionally different from ExternalPlannerServer (8080) + +# โ”€โ”€โ”€ Cameras โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# CAM_KEY = Web camera (opencv, index_or_path=1, /dev/video1) +# CAM2_KEY = RealSense D435i (serial 327743060231) +CAM_KEY = "cam" # Wrist camera โ€” opencv /dev/video6 +# 64x48 preserves the native 640x480 (4:3) capture aspect ratio exactly (10x isotropic +# downsample) โ€” unlike 84x84, which stretches x/y by different factors (84/640 vs 84/480). +CAM_H = 48 +CAM_W = 64 +CAM_C = 3 + +CAM2_KEY = "cam2" # RealSense D435i โ€” /dev/video0, serial 317422074482 +CAM2_H = 48 +CAM2_W = 64 +CAM2_C = 3 diff --git a/mpail2/envs/real/so101/so101_env.py b/mpail2/envs/real/so101/so101_env.py new file mode 100644 index 0000000..bfae4f3 --- /dev/null +++ b/mpail2/envs/real/so101/so101_env.py @@ -0,0 +1,428 @@ +"""SO-101 (SO-ARM101) base Gymnasium environment. + +Communicates with the SO-101 robot-control server (so101_robot_server.py, run in +the `lerobot` conda env) over gRPC. The server directly wraps lerobot's SOFollower +driver and exposes two RPCs (see transport/so101_robot.proto): + + Reset (ResetRequest) -> RobotState + Move the arm to HOME_POSITION_DEG and open the gripper. + + Step (StepRequest{joints_deg}) -> RobotState + Execute a joint-position command and return the resulting state. + +RobotState carries joints_deg (6 floats) plus both camera frames (raw uint8 +bytes) in one message, so each step()/reset() call is a single round-trip +instead of four separate HTTP requests (state + step + camera + camera2). + +Start the server (separate `lerobot` conda env, real hardware attached there): + conda activate lerobot + python so101_robot_server.py --robot_port /dev/ttyACM0 --robot_id Kid \\ + --cam_index /dev/video0 --cam2_serial --grpc_port 7070 +""" + +from __future__ import annotations + +import logging +import time +from typing import Any, Dict, Optional, Tuple + +import gymnasium as gym +import numpy as np +from gymnasium import spaces + +logger = logging.getLogger("so101_env") + +from .robot_limits import ( + ACTION_DIM, + CAM_C, + CAM_H, + CAM_KEY, + CAM_W, + CAM2_C, + CAM2_H, + CAM2_KEY, + CAM2_W, + CONTROL_FREQUENCY_HZ, + DEFAULT_ROBOT_HOST, + DEFAULT_ROBOT_PORT, + EE_LOWER_M, + EE_UPPER_M, + GRIPPER_HALF_RANGE, + HOME_POSITION_DEG, + JOINT_LOWER_DEG, + JOINT_NAMES, + JOINT_UPPER_DEG, + MAX_EPISODE_STEPS, + WRIST_ROLL_HALF_RANGE, + EE_PROPRIO_DIM, +) +from .ik_utils import fk as _fk, ik as _ik, joints_to_ee_proprio as _joints_to_ee_proprio + +try: + import grpc + from .transport import so101_robot_pb2, so101_robot_pb2_grpc + _GRPC_AVAILABLE = True +except ImportError: + _GRPC_AVAILABLE = False + print("[WARNING] grpc / transport.so101_robot_pb2 not available โ€” SO101RobotEnv will run in mock mode.") + + +# obs key used by planner_server.py (--obs_key default) +OBS_KEY = "observation.state" + +# IK sanity-check threshold (degrees of joint movement per millimetre of requested +# Cartesian delta) โ€” see the warning logged in _action_norm_to_joints. A well-conditioned +# solve in this workspace typically lands around 0.03-0.07 deg/mm; this is set well above +# that so only genuinely disproportionate jumps (near-singularity/boundary behavior) fire. +IK_SANITY_DEG_PER_MM = 0.3 + + +class SO101RobotEnv(gym.Env): + """Gymnasium env that wraps the SO-101 arm via the so101_robot_server.py gRPC service. + + ``observation_space`` and ``action_space`` use a leading batch dimension + of 1 so shapes match MPAIL2Runner expectations directly. + """ + + metadata = {"render_modes": []} + + def __init__( + self, + host: str = DEFAULT_ROBOT_HOST, + port: int = DEFAULT_ROBOT_PORT, + control_frequency: float = CONTROL_FREQUENCY_HZ, + max_episode_steps: int = MAX_EPISODE_STEPS, + mock: bool = False, + speed_scale: float = 1.0, + lpf_alpha: float = 1.0, + reset_pause_seconds: float = 3.0, + gripper_hold_steps: int = 1, + ): + super().__init__() + + self.host = host + self.port = port + self.control_frequency = control_frequency + self._step_dt = 1.0 / control_frequency + self.max_episode_steps = max_episode_steps + self.reset_pause_seconds = reset_pause_seconds + # Alias: SO101RealWrapper reads env.unwrapped.max_episode_length (not + # max_episode_steps) to size MPAIL2Runner's rollout length. Without this, + # the wrapper's hasattr() check misses this attribute entirely and silently + # falls back to the MAX_EPISODE_STEPS constant, ignoring whatever value was + # actually requested (e.g. via --max_episode_steps). + self.max_episode_length = max_episode_steps + self.mock = mock or not _GRPC_AVAILABLE + self.speed_scale = speed_scale + # EMA weight on the new action (lower = smoother, more lag). Same role as + # grpc_policy_server.py's --lpf_alpha: an untrained/highly-exploratory MPPI + # policy's raw output is tanh-squashed noise that saturates near +-1 on most + # dimensions most of the time, so without smoothing every step jerks near the + # full (speed_scale-scaled) envelope โ€” reducing speed_scale alone only shrinks + # that envelope, it doesn't stop the action from constantly sitting at its edge. + self._lpf_alpha = lpf_alpha + self._prev_action_norm = None + self._step_count = 0 + self._consecutive_rpc_failures = 0 + # Gripper is excluded from the LPF above (see step()) so it always tracks the raw + # commanded value with no smoothing lag โ€” fine for a fast/light actuator, but it + # also means every step's fresh MPPI noise reaches the gripper directly. Holding + # the commanded gripper position fixed for gripper_hold_steps-1 out of every + # gripper_hold_steps steps turns that into a deliberate, chunked open/close + # command instead of continuous per-step jitter. 1 = update every step (default, + # unchanged behavior). + assert gripper_hold_steps >= 1, "gripper_hold_steps must be >= 1" + self.gripper_hold_steps = gripper_hold_steps + + # runner.py reads these from env.unwrapped + self.num_envs = 1 + + self.observation_space = spaces.Dict({ + OBS_KEY: spaces.Box( + low=-np.inf, high=np.inf, + shape=(1, EE_PROPRIO_DIM), + dtype=np.float32, + ), + CAM_KEY: spaces.Box( + low=0.0, high=1.0, + shape=(1, CAM_H, CAM_W, CAM_C), + dtype=np.float32, + ), + CAM2_KEY: spaces.Box( + low=0.0, high=1.0, + shape=(1, CAM2_H, CAM2_W, CAM2_C), + dtype=np.float32, + ), + }) + # Actions are normalised joint targets in [-1, 1] (planner outputs tanh) + self.action_space = spaces.Box( + low=-1.0, high=1.0, shape=(1, ACTION_DIM), dtype=np.float32 + ) + + self._current_joints = HOME_POSITION_DEG.copy() + self._channel = None + self._stub = None + + if self.mock: + print("[SO101RobotEnv] Running in mock mode โ€” no gRPC calls will be made.") + else: + print(f"[SO101RobotEnv] Connecting to robot server at {host}:{port} (gRPC)") + self._channel = grpc.insecure_channel(f"{host}:{port}") + self._stub = so101_robot_pb2_grpc.SO101RobotStub(self._channel) + + _mock_cam: np.ndarray = None # lazily allocated in mock mode + _mock_cam2: np.ndarray = None + + # โ”€โ”€โ”€ gRPC helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def _decode_state(self, state: "so101_robot_pb2.RobotState") -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """RobotState -> (joints_deg (6,), cam (CAM_H,CAM_W,CAM_C) [0,1], cam2 (CAM2_H,CAM2_W,CAM2_C) [0,1]).""" + joints = np.array(state.joints_deg, dtype=np.float32) + cam = ( + np.frombuffer(state.cam, dtype=np.uint8).reshape(CAM_H, CAM_W, CAM_C).astype(np.float32) / 255.0 + if state.cam else np.zeros((CAM_H, CAM_W, CAM_C), dtype=np.float32) + ) + cam2 = ( + np.frombuffer(state.cam2, dtype=np.uint8).reshape(CAM2_H, CAM2_W, CAM2_C).astype(np.float32) / 255.0 + if state.cam2 else np.zeros((CAM2_H, CAM2_W, CAM2_C), dtype=np.float32) + ) + return joints, cam, cam2 + + def _mock_state(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + if self._mock_cam is None: + self._mock_cam = np.zeros((CAM_H, CAM_W, CAM_C), dtype=np.float32) + if self._mock_cam2 is None: + self._mock_cam2 = np.zeros((CAM2_H, CAM2_W, CAM2_C), dtype=np.float32) + return self._current_joints.copy(), self._mock_cam, self._mock_cam2 + + def _rpc_reset(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, bool]: + if self.mock: + self._current_joints = HOME_POSITION_DEG.copy() + return self._mock_state() + (True,) + try: + state = self._stub.Reset(so101_robot_pb2.ResetRequest(), timeout=8.0) + return self._decode_state(state) + (True,) + except grpc.RpcError as exc: + logger.warning(f"Reset RPC failed ({exc.code()}: {exc.details()}) โ€” continuing with last known state") + return self._mock_state() + (False,) + + def _rpc_step(self, target_deg: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, bool]: + if self.mock: + self._current_joints = target_deg.copy() + return self._mock_state() + (True,) + try: + # 5s, not 2s: a real Step call round-trips send_action() + get_observation() + # (both cameras) server-side; a slow camera/servo read shouldn't be mistaken + # for a dead robot and silently fall back to stale state every subsequent step. + state = self._stub.Step( + so101_robot_pb2.StepRequest(joints_deg=target_deg.tolist()), timeout=5.0 + ) + return self._decode_state(state) + (True,) + except grpc.RpcError as exc: + logger.warning(f"Step RPC failed ({exc.code()}: {exc.details()}) โ€” returning last known state") + return self._mock_state() + (False,) + + # โ”€โ”€โ”€ action scaling โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # NOTE: this env's action space is [x, y, z, wrist_roll, gripper] with xyz as an + # ABSOLUTE position (see _action_norm_to_joints below) but wrist_roll/gripper still + # delta-style. grpc_policy_server.py's own _project_action_to_bounds / + # _action_norm_to_joints were NOT updated to match โ€” it still uses the older + # all-delta Cartesian scheme. A checkpoint trained through one is not + # action-space-compatible with the other. + + def _project_action_to_bounds( + self, + action_norm: np.ndarray, + current_arm_deg: np.ndarray, + current_gripper_deg: float, + ) -> np.ndarray: + """Zero out wrist_roll/gripper components that would push past a hard boundary. + + xyz (indices 0-2) is an ABSOLUTE position command, not a delta, so there's no + "pushing against a wall" concept for it โ€” clipping the mapped target into the + workspace box in _action_norm_to_joints is sufficient. + """ + projected = action_norm.copy() + + wr = current_arm_deg[4] + if wr <= float(JOINT_LOWER_DEG[4]) and projected[3] < 0.0: + projected[3] = 0.0 + elif wr >= float(JOINT_UPPER_DEG[4]) and projected[3] > 0.0: + projected[3] = 0.0 + + g = current_gripper_deg + if g <= float(JOINT_LOWER_DEG[5]) and projected[4] < 0.0: + projected[4] = 0.0 + elif g >= float(JOINT_UPPER_DEG[5]) and projected[4] > 0.0: + projected[4] = 0.0 + + return projected + + def _action_norm_to_joints(self, action_norm: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Decode a 5-dim normalised action [x,y,z,wrist_roll,gripper] to 6 joint degrees via IK. + + x,y,z is an ABSOLUTE target position: action_norm[:3] in [-1,1] maps linearly + onto the EE_LOWER_M/EE_UPPER_M workspace box (matches Kinova's Pick-and-Place + action space in the MPAIL2 paper โ€” position commands, not a delta). wrist_roll + and gripper are still delta/increment style, same as before. + + Returns (joints_deg (6,), effective_action_norm (5,)) โ€” effective_action_norm is the + post-boundary-projection action, used by step() as the LPF reference for next step. + """ + current_arm_deg = self._current_joints[:5].copy() + current_ee = _fk(current_arm_deg) + current_gripper_deg = float(self._current_joints[5]) + + action_norm = self._project_action_to_bounds(action_norm, current_arm_deg, current_gripper_deg) + + target_xyz = EE_LOWER_M + (action_norm[:3].astype(np.float32) + 1.0) / 2.0 * (EE_UPPER_M - EE_LOWER_M) + target_xyz = np.clip(target_xyz, EE_LOWER_M, EE_UPPER_M) + + arm_deg = _ik(target_xyz, initial_arm_deg=current_arm_deg) + arm_deg = np.clip(arm_deg, JOINT_LOWER_DEG[:5], JOINT_UPPER_DEG[:5]) + + # IK sanity check: a small requested Cartesian delta shouldn't produce a + # disproportionately large joint-angle jump. ik()'s DLS solve is redundant + # (3D position target, 4 DOF) with no explicit preference for staying close to + # the previous joint solution beyond the initial_arm_deg warm-start, so near a + # singularity or workspace-boundary configuration its chosen solution can shift + # a lot even for a tiny requested Cartesian delta โ€” which then shows up + # downstream as a target the servo can't track, easily mistaken for a motor/ + # settle problem when the actual target itself was the unstable part. + requested_delta_mm = float(np.linalg.norm(target_xyz - current_ee)) * 1000.0 + joint_delta_deg = arm_deg - current_arm_deg + worst_idx = int(np.argmax(np.abs(joint_delta_deg))) + max_joint_delta_deg = float(np.abs(joint_delta_deg[worst_idx])) + if requested_delta_mm > 1e-3: + deg_per_mm = max_joint_delta_deg / requested_delta_mm + if deg_per_mm > IK_SANITY_DEG_PER_MM: + logger.warning( + f"[ik] disproportionate joint jump: requested {requested_delta_mm:.2f}mm Cartesian " + f"delta -> {max_joint_delta_deg:.2f}ยฐ on {JOINT_NAMES[worst_idx]} ({deg_per_mm:.2f}ยฐ/mm) " + f"at EE={tuple(np.round(current_ee, 4).tolist())} target_xyz={tuple(np.round(target_xyz, 4).tolist())}" + ) + + # wrist_roll is back as an action component (index 3), delta-style โ€” ik() itself + # still never touches index 4 of arm_deg (see its docstring), so we set it here. + wrist_delta_deg = float(action_norm[3]) * float(WRIST_ROLL_HALF_RANGE) * self.speed_scale + arm_deg[4] = float(np.clip(current_arm_deg[4] + wrist_delta_deg, + JOINT_LOWER_DEG[4], JOINT_UPPER_DEG[4])) + + # Only accept a new gripper command every gripper_hold_steps'th step (starting at + # step 0 of each episode); on held steps the target stays exactly where it is, and + # the effective action is zeroed to match (so logged/replayed action_executed + # reflects what was actually applied, not what was requested-but-held). + gripper_update_due = (self._step_count % self.gripper_hold_steps) == 0 + if not gripper_update_due: + action_norm[4] = 0.0 + gripper_delta_deg = float(action_norm[4]) * float(GRIPPER_HALF_RANGE) * self.speed_scale + gripper_deg = float(np.clip(current_gripper_deg + gripper_delta_deg, + JOINT_LOWER_DEG[5], JOINT_UPPER_DEG[5])) + + return np.append(arm_deg, gripper_deg).astype(np.float32), action_norm + + # โ”€โ”€โ”€ Gymnasium API โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def _make_obs(self, joints: np.ndarray, cam: np.ndarray, cam2: np.ndarray) -> Dict[str, np.ndarray]: + ee_proprio = _joints_to_ee_proprio(joints) + return { + OBS_KEY: ee_proprio.reshape(1, EE_PROPRIO_DIM).astype(np.float32), + CAM_KEY: cam.reshape(1, CAM_H, CAM_W, CAM_C), + CAM2_KEY: cam2.reshape(1, CAM2_H, CAM2_W, CAM2_C), + } + + def reset( + self, + *, + seed: Optional[int] = None, + options: Optional[Dict[str, Any]] = None, + ) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]: + super().reset(seed=seed) + t0 = time.time() + self._prev_action_norm = None # don't carry LPF state across episode boundaries + self._step_count = 0 + self._consecutive_rpc_failures = 0 + self._current_joints, cam, cam2, ok = self._rpc_reset() + logger.info( + f"[reset] ok={ok} joints_deg={np.round(self._current_joints, 2).tolist()} " + f"took={time.time() - t0:.2f}s" + ) + if self.reset_pause_seconds > 0: + logger.info(f"[reset] pausing {self.reset_pause_seconds:.1f}s for the arm to settle at home...") + time.sleep(self.reset_pause_seconds) + info: Dict[str, Any] = { + "reset_time": round(time.time() - t0, 3), + "joint_names": JOINT_NAMES, + } + return self._make_obs(self._current_joints, cam, cam2), info + + def step( + self, action: np.ndarray + ) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]: + t0 = time.time() + + raw_action = np.asarray(action, dtype=np.float32).reshape(-1) + action_vec = raw_action.copy() + + # LPF on xyz + wrist_roll only (indices 0-3); gripper (index 4) is excluded so it + # always tracks the raw commanded value, same convention as grpc_policy_server.py. + if self._prev_action_norm is not None: + gripper_val = action_vec[4] + action_vec = (self._lpf_alpha * action_vec + + (1 - self._lpf_alpha) * self._prev_action_norm) + action_vec[4] = gripper_val + action_vec = np.clip(action_vec, -1.0, 1.0) + + prev_joints = self._current_joints.copy() + target_deg, effective_norm = self._action_norm_to_joints(action_vec) + self._prev_action_norm = effective_norm.copy() + self._current_joints, cam, cam2, ok = self._rpc_step(target_deg) + obs = self._make_obs(self._current_joints, cam, cam2) + + if ok: + self._consecutive_rpc_failures = 0 + else: + self._consecutive_rpc_failures += 1 + if self._consecutive_rpc_failures >= 3: + logger.error( + f"[step {self._step_count}] Step RPC has failed " + f"{self._consecutive_rpc_failures} times in a row โ€” joints have not " + f"actually updated since. Check so101_robot_server.py (camera/servo hang?) " + f"or raise the RPC timeout further if this is just slow, not stuck." + ) + + ee_xyz = _fk(self._current_joints[:5]) + moved = self._current_joints - prev_joints + logger.info( + f"[step {self._step_count}] ok={ok} raw={np.round(raw_action, 3).tolist()} " + f"eff={np.round(effective_norm, 3).tolist()} " + f"target_deg={np.round(target_deg, 2).tolist()} " + f"actual_deg={np.round(self._current_joints, 2).tolist()} " + f"delta_deg={np.round(moved, 2).tolist()} " + f"EE=({ee_xyz[0]:+.4f},{ee_xyz[1]:+.4f},{ee_xyz[2]:+.4f}) " + f"elapsed={time.time() - t0:.3f}s" + ) + self._step_count += 1 + + elapsed = time.time() - t0 + remaining = self._step_dt - elapsed + if remaining > 0: + time.sleep(remaining) + + info: Dict[str, Any] = { + "joint_positions_deg": self._current_joints.tolist(), + "action_norm": action_vec.tolist(), + "rpc_ok": ok, + "mpail_env/step_time": round(time.time() - t0, 3), + } + # Reward is always 0 โ€” MPAIL learns reward purely from demonstrations + return obs, 0.0, False, False, info + + def read_robot_state(self) -> np.ndarray: + """Return current joint positions in degrees (shape: (STATE_DIM,)).""" + return self._current_joints.copy() + + def close(self) -> None: + if self._channel is not None: + self._channel.close() + super().close() diff --git a/mpail2/envs/real/so101/soa.urdf b/mpail2/envs/real/so101/soa.urdf new file mode 100644 index 0000000..ba41bf3 --- /dev/null +++ b/mpail2/envs/real/so101/soa.urdf @@ -0,0 +1,386 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mpail2/envs/real/so101/training/__init__.py b/mpail2/envs/real/so101/training/__init__.py new file mode 100644 index 0000000..0af8045 --- /dev/null +++ b/mpail2/envs/real/so101/training/__init__.py @@ -0,0 +1 @@ +"""Standalone SO-101 training/data entry points (demo recording, conversion, training, diagnostics).""" diff --git a/mpail2/envs/real/so101/training/check_encoder_collapse.py b/mpail2/envs/real/so101/training/check_encoder_collapse.py new file mode 100644 index 0000000..697f811 --- /dev/null +++ b/mpail2/envs/real/so101/training/check_encoder_collapse.py @@ -0,0 +1,165 @@ +"""check_encoder_collapse.py โ€” Diagnose whether the encoder's latent space has +collapsed (encoding most/all observations to a low-variance region), which a +low dynamics (JEP) loss alone cannot rule out: an encoder that maps everything +to a near-constant vector also gets near-zero prediction loss, without having +learned anything useful. + +Loads a checkpoint + demo.pt, runs the encoder over a sample of real (obs, +next_obs) pairs (through the same CamOnlyObsNormalizer pipeline used during +training), and reports: + - per-dimension std (how many latent dims are effectively "dead") + - effective rank of the latent covariance (participation ratio + cumulative + explained-variance curve): a healthy, non-collapsed representation should + spread variance across many dimensions, not concentrate it in a handful. + +Usage: + python check_encoder_collapse.py --demo_path demo.pt \\ + --checkpoint logs/so101_local/models/model_65.pt +""" + +import argparse +import glob +import os + +import numpy as np +import torch + +from mpail2.envs.real.so101 import ik_utils +from mpail2.envs.real.so101 import OBS_KEY, EE_PROPRIO_DIM +from mpail2.envs.real.so101.robot_limits import ( + CAM_KEY, CAM_H, CAM_W, CAM_C, CAM2_KEY, CAM2_H, CAM2_W, CAM2_C, +) +from mpail2.configs.defs import MultiCoderConfig, CNNCoderConfig, PlannerConfig, PolicySamplingConfig +from mpail2.planner import Planner +from mpail2.utils.obs_normalizer import ObsNormalizerFactory + + +def latest_checkpoint(models_dir: str) -> str: + candidates = glob.glob(os.path.join(models_dir, "model_*.pt")) + if not candidates: + raise FileNotFoundError(f"No checkpoints found under {models_dir}") + return max(candidates, key=os.path.getmtime) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--demo_path", default="demo.pt") + parser.add_argument("--checkpoint", default=None, + help="Path to a model_*.pt checkpoint. Defaults to the most " + "recently modified one under logs/so101_local/models.") + parser.add_argument("--latent_dim", type=int, default=512) + parser.add_argument("--joint_dim", type=int, default=256) + parser.add_argument("--no_wrist_cam", action="store_true") + parser.add_argument("--num_pairs", type=int, default=4096, + help="Number of (obs, next_obs) demo transitions to sample " + "(yields 2x this many latent vectors).") + parser.add_argument("--device", default="cuda") + parser.add_argument("--random_init", action="store_true", + help="Skip checkpoint loading, use freshly initialized weights " + "(architecture-only sanity check, not a trained-model diagnosis).") + args = parser.parse_args() + + device = args.device if (args.device != "cuda" or torch.cuda.is_available()) else "cpu" + checkpoint = None if args.random_init else (args.checkpoint or latest_checkpoint("logs/so101_local/models")) + print(f"Checkpoint: {checkpoint if checkpoint else '(random init โ€” no checkpoint loaded)'}") + print(f"Demo file : {args.demo_path}") + + demonstrations = torch.load(args.demo_path, map_location="cpu", weights_only=False) + demo_keys = {OBS_KEY, CAM_KEY, CAM2_KEY} + demonstrations = {k: v.float() for k, v in demonstrations.items() if k in demo_keys} + if CAM2_KEY not in demonstrations: + demonstrations[CAM2_KEY] = torch.zeros( + demonstrations[CAM_KEY].shape[:-1] + (CAM2_C,), dtype=torch.float32 + ) + + print("Converting demo joint states to EE proprioception (running FK)...") + demo_joints = demonstrations[OBS_KEY].numpy() + N = demo_joints.shape[0] + ee_proprio = np.zeros((N, 2, EE_PROPRIO_DIM), dtype=np.float32) + for i in range(N): + for j in range(2): + ee_proprio[i, j] = ik_utils.joints_to_ee_proprio(demo_joints[i, j]) + demonstrations[OBS_KEY] = torch.from_numpy(ee_proprio) + + num_pairs = min(args.num_pairs, N) + idx = np.random.choice(N, size=num_pairs, replace=False) + + _coder_list = [ + MultiCoderConfig.ProprioCoderConfig(obs_key=OBS_KEY, input_dim=EE_PROPRIO_DIM, output_dim=args.joint_dim), + CNNCoderConfig(obs_key=CAM2_KEY, H=CAM2_H, W=CAM2_W, C=CAM2_C), + ] + if not args.no_wrist_cam: + _coder_list.append(CNNCoderConfig(obs_key=CAM_KEY, H=CAM_H, W=CAM_W, C=CAM_C)) + encoder_cfg = MultiCoderConfig(coder_list=_coder_list) + + planner_cfg = PlannerConfig( + encoder_cfg=encoder_cfg, action_dim=5, latent_dim=args.latent_dim, + sampling_cfg=PolicySamplingConfig(), + ) + + planner = Planner(policy_config=planner_cfg, num_envs=1, device=device, dtype=torch.float32) + if checkpoint is not None: + saved_dict = torch.load(checkpoint, map_location=device, weights_only=False) + planner.load_state_dict(saved_dict["model_state_dict"]) + planner.eval() + + obs_normalizer = ObsNormalizerFactory.create_normalizer(normalization_type="cam_only", device=device) + + # Stack obs (t) and next_obs (t+1) together for more diverse samples. + obs = { + OBS_KEY: torch.cat([demonstrations[OBS_KEY][idx, 0], demonstrations[OBS_KEY][idx, 1]], dim=0).to(device), + CAM_KEY: torch.cat([demonstrations[CAM_KEY][idx, 0], demonstrations[CAM_KEY][idx, 1]], dim=0).to(device), + CAM2_KEY: torch.cat([demonstrations[CAM2_KEY][idx, 0], demonstrations[CAM2_KEY][idx, 1]], dim=0).to(device), + } + obs = obs_normalizer(obs) + + latents = [] + batch_size = 512 + total = obs[OBS_KEY].shape[0] + with torch.no_grad(): + for start in range(0, total, batch_size): + end = min(start + batch_size, total) + batch = {k: v[start:end] for k, v in obs.items()} + latents.append(planner.encoder(batch).cpu()) + Z = torch.cat(latents, dim=0) # [2*num_pairs, latent_dim] + + print(f"\nEncoded {Z.shape[0]} latent vectors, dim={Z.shape[1]}") + + mean = Z.mean(dim=0) + std = Z.std(dim=0) + print(f"\nPer-dim latent std: min={std.min():.4g} median={std.median():.4g} " + f"mean={std.mean():.4g} max={std.max():.4g}") + + for thresh_frac in (0.01, 0.05, 0.10): + thresh = thresh_frac * std.max() + n_dead = (std < thresh).sum().item() + print(f" dims with std < {thresh_frac*100:.0f}% of max std ({thresh:.4g}): " + f"{n_dead} / {std.numel()} ({100*n_dead/std.numel():.1f}%)") + + centered = Z - mean + # SVD of centered latents for the variance spectrum (avoids materializing the + # full [dim, dim] covariance matrix). + _, S, _ = torch.linalg.svd(centered, full_matrices=False) + eigvals = (S ** 2) / (Z.shape[0] - 1) + total_var = eigvals.sum() + explained = torch.cumsum(eigvals, dim=0) / total_var + + participation_ratio = (eigvals.sum() ** 2) / (eigvals ** 2).sum() + print(f"\nEffective rank (participation ratio): {participation_ratio:.1f} / {Z.shape[1]} dims") + + for frac in (0.50, 0.90, 0.99): + n_dims = int((explained >= frac).nonzero()[0].item()) + 1 + print(f" dims needed for {frac*100:.0f}% of variance: {n_dims} / {Z.shape[1]}") + + print(f"\nTop-5 eigenvalue share of total variance: " + f"{[f'{(e/total_var).item():.3f}' for e in eigvals[:5]]}") + + print("\n--- Interpretation ---") + print(f"latent_dim={Z.shape[1]}. If effective rank is a small fraction of latent_dim, " + f"or a handful of dims explain >90% of variance, or many dims have near-zero std, " + f"that's the signature of representation collapse (encoder ignoring most of what " + f"it's given, dynamics loss looking good for the wrong reason).") + + +if __name__ == "__main__": + main() diff --git a/mpail2/envs/real/so101/training/convert.py b/mpail2/envs/real/so101/training/convert.py new file mode 100644 index 0000000..3d827d4 --- /dev/null +++ b/mpail2/envs/real/so101/training/convert.py @@ -0,0 +1,126 @@ +"""Convert raw_demos*/*.npz trajectory files to a single MPAIL-format .pt file. + +Usage: + python convert.py # merges raw_demos2 only (default) + python convert.py --dirs raw_demos2 # same + python convert.py --dirs raw_demos raw_demos2 raw_demos3 # all demos + python convert.py --out my_demos.pt # custom output path + python convert.py --img_w 64 --img_h 48 # resize cameras to 64x48 (default, matches 640x480 aspect ratio) +""" + +import argparse +from pathlib import Path + +import cv2 +import numpy as np +import torch + +_IMG_KEY_PREFIXES = ("observation.images.",) + + +def _normalize_key(key: str) -> str: + """Strip lerobot image key prefix: 'observation.images.cam' โ†’ 'cam'.""" + for prefix in _IMG_KEY_PREFIXES: + if key.startswith(prefix): + return key[len(prefix):] + return key + + +def _resize_cam(tensor: torch.Tensor, out_w: int, out_h: int) -> torch.Tensor: + """Resize camera tensor (N, 2, H, W, C) โ†’ (N, 2, out_h, out_w, C). + + Uses cv2.INTER_AREA to match the downsampling used for live camera frames + in so101_robot_server.py โ€” keeping demo and online observations consistent. + out_w/out_h should match the source aspect ratio (e.g. 64x48 for 640x480 + native capture) to avoid the anisotropic stretch a square resize would cause. + """ + N, two, H, W, C = tensor.shape + if H == out_h and W == out_w: + return tensor + arr = tensor.reshape(N * two, H, W, C).numpy() + if arr.max() > 1.0: + arr = arr / 255.0 + resized = np.stack([ + cv2.resize(frame, (out_w, out_h), interpolation=cv2.INTER_AREA) + for frame in arr + ]) + return torch.from_numpy(resized).reshape(N, two, out_h, out_w, C) + + +def convert(source_dirs: list[str], out_path: str, img_w: int = 64, img_h: int = 48, state_dim: int | None = None) -> None: + all_data: dict = {} + total_files = 0 + + for dir_name in source_dirs: + src = Path(dir_name) + if not src.exists(): + print(f"[WARN] {src} does not exist โ€” skipping") + continue + npz_files = sorted(src.glob("*.npz")) + if not npz_files: + print(f"[WARN] {src} contains no .npz files โ€” skipping") + continue + for npz_file in npz_files: + d = np.load(str(npz_file), allow_pickle=True) + for raw_key in d.files: + key = _normalize_key(raw_key) + arr = torch.from_numpy(d[raw_key].astype(np.float32)) + # Resize cameras immediately to avoid holding 480ร—640 frames in RAM + if arr.dim() == 5: + arr = _resize_cam(arr, img_w, img_h) + all_data.setdefault(key, []).append(arr) + total_files += 1 + print(f" loaded {npz_file} keys={list(d.files)}") + + if not all_data: + raise RuntimeError("No data found in any of the source directories.") + + merged = {k: torch.cat(vs, dim=0) for k, vs in all_data.items()} + + # Keep only MPAIL-format keys: shape (N, 2, *obs_shape) โ€” drop "actions" etc. + demos = {k: v for k, v in merged.items() if v.dim() >= 3 and v.shape[1] == 2} + + if not demos: + raise RuntimeError( + "No MPAIL-format tensors found (expected shape [N, 2, *obs_shape]). " + "Keys found: " + str(list(merged.keys())) + ) + + # Trim state vector if it has extra dims (e.g. follower+leader concatenated โ†’ follower only) + state_key = "observation.state" + if state_dim is not None and state_key in demos and demos[state_key].shape[-1] != state_dim: + before = demos[state_key].shape[-1] + demos[state_key] = demos[state_key][..., :state_dim] + print(f" trimmed {state_key}: dim {before} โ†’ {state_dim}") + + torch.save(demos, out_path) + print(f"\nSaved {out_path}") + print(f" Source files : {total_files}") + for k, v in demos.items(): + print(f" {k}: {tuple(v.shape)} ({v.numel() * 4 / 1e6:.2f} MB)") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dirs", nargs="+", default=["raw_demos2"], + help="Source directories containing .npz trajectory files (default: raw_demos2)", + ) + parser.add_argument( + "--out", default="raw_demos2_master.pt", + help="Output .pt file path (default: raw_demos2_master.pt)", + ) + parser.add_argument( + "--img_w", type=int, default=64, + help="Resize camera images to this width (default: 64 โ€” matches 640x480 native aspect ratio).", + ) + parser.add_argument( + "--img_h", type=int, default=48, + help="Resize camera images to this height (default: 48 โ€” matches 640x480 native aspect ratio).", + ) + parser.add_argument( + "--state_dim", type=int, default=None, + help="Trim observation.state to this many dims (e.g. 6 if data has follower+leader=12)", + ) + args = parser.parse_args() + convert(args.dirs, args.out, img_w=args.img_w, img_h=args.img_h, state_dim=args.state_dim) diff --git a/mpail2/envs/real/so101/training/convert_lerobot.py b/mpail2/envs/real/so101/training/convert_lerobot.py new file mode 100644 index 0000000..c97c45c --- /dev/null +++ b/mpail2/envs/real/so101/training/convert_lerobot.py @@ -0,0 +1,141 @@ +""" +convert_lerobot.py โ€” Convert a LeRobot recorded dataset to MPAIL2 .pt format. + +Run in the LEROBOT conda env (lerobot must be installed): + conda activate lerobot + python convert_lerobot.py \ + --repo_id local/my_demos \ + --root ./raw_lerobot_demos \ + --out raw_demos2_master.pt \ + --cam_name cam \ + --img_w 64 --img_h 48 + +What it does +โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + For every consecutive frame pair (t, t+1) within each episode it builds: + observation.state : (N, 2, STATE_DIM) joint positions + cam : (N, 2, 48, 64, 3) camera frames, resized & in [0, 1] + + N = total transitions across all episodes (last frame of each episode is dropped + because it has no t+1 counterpart). + +The output file is ready for: + python train_so101_local.py --demo_path raw_demos2_master.pt +""" + +import argparse +from pathlib import Path + +import cv2 +import numpy as np +import torch + + +def convert(repo_id: str, root: str | None, out: str, cam_name: str, img_w: int, img_h: int): + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + print(f"Loading dataset repo_id={repo_id} root={root or 'hub'} ...") + ds = LeRobotDataset(repo_id=repo_id, root=root) + print(f" {ds.num_episodes} episodes, {ds.num_frames} frames total") + print(f" features: {list(ds.features.keys())}") + + lerobot_cam_key = f"observation.images.{cam_name}" + has_camera = lerobot_cam_key in ds.features + if not has_camera: + print(f"[WARN] Camera key '{lerobot_cam_key}' not found. Available image keys:") + for k in ds.features: + if "image" in k: + print(f" {k}") + print(" Continuing without camera data.") + + state_pairs = [] + cam_pairs = [] + + # Iterate episode by episode so we never pair frames across boundaries + hf = ds.hf_dataset + episode_indices = sorted(hf.unique("episode_index")) + print(f"Converting {len(episode_indices)} episodes ...") + + for ep_idx in episode_indices: + ep_rows = hf.filter(lambda row: row["episode_index"] == ep_idx) + n_frames = len(ep_rows) + if n_frames < 2: + print(f" episode {ep_idx}: only {n_frames} frame(s), skipping") + continue + + # Pull state โ€” shape (T, STATE_DIM) + states = np.array(ep_rows["observation.state"], dtype=np.float32) # (T, D) + + # Build consecutive pairs (T-1, 2, D) + state_t = states[:-1] # (T-1, D) + state_t1 = states[1:] # (T-1, D) + state_pairs.append(np.stack([state_t, state_t1], axis=1)) # (T-1, 2, D) + + if has_camera: + # Frames are decoded tensors; lerobot returns (C, H, W) uint8 or float + cam_frames = [] + for i in range(n_frames): + frame = ds[ep_rows[i]["frame_index"]] + img = frame[lerobot_cam_key] # tensor (C, H, W) + if img.dtype == torch.uint8: + img = img.float() / 255.0 + cam_frames.append(img) # (C, H, W) float [0,1] + + # Stack โ†’ (T, C, H, W), resize (cv2.INTER_AREA, matching live camera + # downsampling in so101_robot_server.py), convert to (T, H, W, C) + cam_t = torch.stack(cam_frames[:-1]).permute(0, 2, 3, 1).numpy() # (T-1, H, W, C) + cam_t1 = torch.stack(cam_frames[1:]).permute(0, 2, 3, 1).numpy() # (T-1, H, W, C) + + cam_t = np.stack([ + cv2.resize(img, (img_w, img_h), interpolation=cv2.INTER_AREA) for img in cam_t + ]) + cam_t1 = np.stack([ + cv2.resize(img, (img_w, img_h), interpolation=cv2.INTER_AREA) for img in cam_t1 + ]) + + cam_pairs.append(np.stack([cam_t, cam_t1], axis=1)) # (T-1, 2, H, W, C) + + print(f" episode {ep_idx}: {n_frames} frames โ†’ {n_frames - 1} transitions") + + # Concatenate across episodes + all_states = np.concatenate(state_pairs, axis=0) # (N, 2, D) + demos = {"observation.state": torch.from_numpy(all_states)} + + if cam_pairs: + all_cams = np.concatenate(cam_pairs, axis=0) # (N, 2, H, W, C) + demos["cam"] = torch.from_numpy(all_cams) + + # Save + torch.save(demos, out) + print(f"\nSaved โ†’ {out}") + for k, v in demos.items(): + mb = v.numel() * 4 / 1e6 + print(f" {k}: {tuple(v.shape)} ({mb:.1f} MB)") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Convert LeRobot dataset โ†’ MPAIL2 .pt") + parser.add_argument("--repo_id", required=True, + help="LeRobot repo ID, e.g. 'local/my_demos'") + parser.add_argument("--root", default=None, + help="Local root directory where the dataset is saved " + "(e.g. ./raw_lerobot_demos). Omit to load from HF Hub.") + parser.add_argument("--out", default="raw_demos2_master.pt", + help="Output .pt file path (default: raw_demos2_master.pt)") + parser.add_argument("--cam_name", default="cam", + help="Camera name used during recording (default: cam). " + "This is the name after 'observation.images.' in the dataset.") + parser.add_argument("--img_w", type=int, default=64, + help="Resize camera frames to this width (default: 64 โ€” matches 640x480 aspect ratio)") + parser.add_argument("--img_h", type=int, default=48, + help="Resize camera frames to this height (default: 48 โ€” matches 640x480 aspect ratio)") + args = parser.parse_args() + + convert( + repo_id=args.repo_id, + root=args.root, + out=args.out, + cam_name=args.cam_name, + img_w=args.img_w, + img_h=args.img_h, + ) diff --git a/mpail2/envs/real/so101/training/demo_recording_server.py b/mpail2/envs/real/so101/training/demo_recording_server.py new file mode 100644 index 0000000..3b56dcb --- /dev/null +++ b/mpail2/envs/real/so101/training/demo_recording_server.py @@ -0,0 +1,443 @@ +""" +demo_recording_server.py โ€” gRPC demo-recording server for SO-101. + +Returns home-position actions so the robot holds still while you move it manually +or via a leader arm. Every (obs_t, obs_t+1) pair is saved to .npz files. + +Run this in the `mpail2` conda env: + conda activate mpail2 + python -m mpail2.envs.real.so101.training.demo_recording_server --collect_dir ./raw_demos2 --flush_every 50 + +After collecting, convert to .pt: + python -m mpail2.envs.real.so101.training.convert --dirs raw_demos2 --out raw_demos2_master.pt + +Robot client (lerobot env) โ€” see mpail2/envs/real/so101/README.md for the full connect + +collect walkthrough and an explanation of each flag: + python -m lerobot.async_inference.robot_client \\ + --robot.type=so100_follower \\ + --robot.port=/dev/ttyACM0 \\ + --robot.id=Kid \\ + --robot.cameras="{cam: {type: opencv, index_or_path: /dev/video0, width: 640, height: 480, fps: 30}, cam2: {type: intelrealsense, serial_number_or_name: 317422074482, width: 640, height: 480, fps: 30}}" \\ + --teleop.type=so100_leader \\ + --teleop.port=/dev/ttyACM1 \\ + --teleop.id=Mom \\ + --server_address=127.0.0.1:8080 \\ + --policy_type=act \\ + --pretrained_name_or_path=dummy \\ + --actions_per_chunk=1 \\ + --task="pick up the cup" +""" + +import argparse +import io +import logging +import pickle # nosec +import threading +import time +from concurrent import futures +from pathlib import Path +from queue import Empty, Queue +from typing import Dict, List, Optional + +import grpc +import numpy as np +import torch + +from mpail2.envs.real.so101.transport import services_pb2, services_pb2_grpc +from lerobot.async_inference.helpers import TimedObservation, TimedAction + +from mpail2.envs.real.so101 import OBS_KEY, HOME_POSITION_DEG + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True) +logger = logging.getLogger("demo_recording_server") + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# CAMERA DISPLAY โ€” live OpenCV window updated from a background thread +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +_CAM_SAVE_PATH = "cameras_latest.png" + +def _update_camera_display(image_arrays: dict) -> None: + """Save both camera images side-by-side to cameras_latest.png. + + Open this file in VSCode โ€” it auto-refreshes as the file updates. + Layout: RealSense (cam2) on the left, wrist cam (cam) on the right. + """ + try: + import cv2 + except ImportError: + return + frames = [] + for key in ("cam2", "cam"): + img = image_arrays.get(key) + if img is None: + continue + img = np.array(img, dtype=np.float32) + if img.max() > 1.0: + img = img / 255.0 + img_u8 = (img * 255).clip(0, 255).astype(np.uint8) + if img_u8.shape[2] == 3: + img_u8 = cv2.cvtColor(img_u8, cv2.COLOR_RGB2BGR) + img_u8 = cv2.resize(img_u8, (320, 240)) + frames.append(img_u8) + if not frames: + return + composed = np.concatenate(frames, axis=1) + cv2.imwrite(_CAM_SAVE_PATH, composed) + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# PICKLE โ€” deserialize lerobot TimedObservation +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +def _loads(data: bytes): + return pickle.loads(data) # nosec + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# OBSERVATION HELPERS +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +def _parse_raw_obs(raw_obs: dict) -> tuple[np.ndarray, dict]: + joint_vals = [] + image_arrays = {} + for key, value in raw_obs.items(): + if key in ("task", "teleop_action"): + continue + arr = np.array(value, dtype=np.float32) + if arr.ndim == 3: + # Camera image (H, W, C); strip lerobot prefix if present + cam_key = key.removeprefix("observation.images.") + image_arrays[cam_key] = arr + else: + # Scalar motor key ("shoulder_pan.pos", ...) or full "observation.state" array + joint_vals.extend(arr.flatten().tolist()) + return np.array(joint_vals, dtype=np.float32), image_arrays + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# DEMO RECORDER (same logic as planner_server.py _record_step / _flush) +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +class DemoRecorder: + """Buffers (obs_t, obs_t+1) pairs and flushes to .npz files.""" + + def __init__(self, collect_dir: Path, flush_every: int = 200): + self.collect_dir = collect_dir + self.flush_every = flush_every + collect_dir.mkdir(parents=True, exist_ok=True) + + self._lock = threading.Lock() + self._obs_buf: List[np.ndarray] = [] + self._next_obs_buf: List[np.ndarray] = [] + self._image_bufs: Dict[str, List[np.ndarray]] = {} + self._image_next_bufs: Dict[str, List[np.ndarray]] = {} + self._pending_obs: Optional[np.ndarray] = None + self._pending_images: Dict[str, np.ndarray] = {} + self._file_idx = 0 + + @staticmethod + def _to_uint8(img: np.ndarray) -> np.ndarray: + if img.max() <= 1.0: + return (img * 255).clip(0, 255).astype(np.uint8) + return img.clip(0, 255).astype(np.uint8) + + def record(self, obs: np.ndarray, images: Dict[str, np.ndarray]): + images_u8 = {k: self._to_uint8(v) for k, v in images.items()} + with self._lock: + if self._pending_obs is not None: + self._obs_buf.append(self._pending_obs) + self._next_obs_buf.append(obs.copy()) + for cam, img in self._pending_images.items(): + self._image_bufs.setdefault(cam, []).append(img) + self._image_next_bufs.setdefault(cam, []).append( + images_u8.get(cam, img).copy() + ) + + self._pending_obs = obs.copy() + self._pending_images = images_u8 + + if len(self._obs_buf) >= self.flush_every: + self._flush() + + def flush(self): + with self._lock: + self._flush() + + def reset_pending(self): + """Discard any half-open transition so the next call to record() starts fresh.""" + with self._lock: + self._pending_obs = None + self._pending_images = {} + + def _flush(self): + if not self._obs_buf: + return + obs_arr = np.stack(self._obs_buf) + nobs_arr = np.stack(self._next_obs_buf) + save_data = {OBS_KEY: np.stack([obs_arr, nobs_arr], axis=1)} # (N, 2, state_dim) + for cam in self._image_bufs: + imgs = np.stack(self._image_bufs[cam]) + nimgs = np.stack(self._image_next_bufs[cam]) + save_data[cam] = np.stack([imgs, nimgs], axis=1) # (N, 2, H, W, C) + + path = self.collect_dir / f"traj_{self._file_idx:04d}.npz" + np.savez(str(path), **save_data) + n = len(self._obs_buf) + cams = list(self._image_bufs.keys()) + logger.info(f"[recorder] Saved {n} steps โ†’ {path} cameras={cams or 'none'}") + + self._obs_buf.clear(); self._next_obs_buf.clear() + self._image_bufs.clear(); self._image_next_bufs.clear() + self._file_idx += 1 + self._pending_obs = None + self._pending_images = {} + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# gRPC SERVICER +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +class DemoRecordingServicer(services_pb2_grpc.AsyncInferenceServicer): + + def __init__( + self, + recorder: DemoRecorder, + max_episode_steps: int = 200, + show_cameras: bool = False, + episode_pause_seconds: float = 5.0, + ): + self.recorder = recorder + self.max_episode_steps = max_episode_steps + self.show_cameras = show_cameras + # Auto-advance episodes by step count instead of waiting for a client-driven + # Ready() call (see _episode_end). + self.episode_pause_seconds = episode_pause_seconds + self._going_home = False + if show_cameras: + logger.info(f"[cameras] Saving live frames to {_CAM_SAVE_PATH} โ€” open in VSCode to watch") + + self._obs_queue: Queue = Queue(maxsize=1) + self.shutdown_event = threading.Event() + + self._episode_steps = 0 + self._episode_count = 0 + self._prev_timestep = -1 + self._recording_paused = False # True while waiting for Enter between episodes + + def Ready(self, request, context): + # Kick off end-of-episode work (save, wait for Enter) in a background thread so + # the client gets this response immediately and the arm can keep moving. + # _recording_paused is set to True inside _episode_end() before any slow work + # starts, so incoming observations are discarded until the next episode is + # ready to record. + if self._episode_steps > 0: + logger.info(f"Ready: starting episode-end work in background (ep {self._episode_count + 1}, {self._episode_steps} steps)...") + threading.Thread(target=self._episode_end, daemon=True).start() + logger.info(f"Robot client connected: {context.peer()}") + self.shutdown_event.clear() + self._obs_queue = Queue(maxsize=1) + return services_pb2.Empty() + + def SendPolicyInstructions(self, request, context): + # No-op โ€” demo recording always returns a single-action chunk regardless of + # the client's requested actions_per_chunk. Implemented only so the RPC isn't + # left unimplemented against lerobot's async_inference client protocol. + return services_pb2.Empty() + + def SendObservations(self, request_iterator, context): + try: + buf = io.BytesIO() + n_chunks = 0 + for chunk in request_iterator: + state = chunk.transfer_state + if state == services_pb2.TransferState.TRANSFER_BEGIN: + buf.seek(0); buf.truncate(0) + buf.write(chunk.data) + n_chunks += 1 + if state == services_pb2.TransferState.TRANSFER_END: + break + data = buf.getvalue() + logger.debug(f"SendObservations: received {n_chunks} chunk(s), {len(data)} bytes") + except Exception as e: + logger.exception(f"SendObservations chunk read error: {e}") + return services_pb2.Empty() + + # While paused between episodes, discard immediately โ€” arm movement is + # handled client-side (leaderโ†’follower), so we don't need to queue or record. + if self._recording_paused: + logger.debug("SendObservations: discarded (recording paused)") + return services_pb2.Empty() + + try: + timed_obs: TimedObservation = _loads(data) + except Exception as e: + logger.exception(f"SendObservations deserialize error: {e} data[:64]={data[:64]!r}") + return services_pb2.Empty() + + try: + timestep = timed_obs.get_timestep() + joint_state, image_arrays = _parse_raw_obs(timed_obs.get_observation()) + + if self._episode_steps == 0: + cam_info = {k: v.shape for k, v in image_arrays.items()} + logger.info(f"[obs check] joints={joint_state.shape} cameras={cam_info}") + + if self.show_cameras: + _update_camera_display(image_arrays) + + if self._episode_steps % 10 == 0: + from mpail2.envs.real.so101.robot_limits import JOINT_NAMES + joint_str = " ".join(f"{n}={v:.2f}" for n, v in zip(JOINT_NAMES, joint_state)) + logger.info(f"[joints] {joint_str}") + + self._prev_timestep = timestep + + self.recorder.record(joint_state, image_arrays) + + # Auto-advance to the next episode once max_episode_steps is reached + # rather than waiting on a client-driven Ready() call. + if not self._recording_paused: + self._episode_steps += 1 + logger.info(f"[ep {self._episode_count + 1} step {self._episode_steps}/{self.max_episode_steps}]") + if self._episode_steps >= self.max_episode_steps: + self._recording_paused = True + threading.Thread(target=self._auto_episode_end, daemon=True).start() + + # Push to obs queue so GetActions can respond + if self._obs_queue.full(): + try: self._obs_queue.get_nowait() + except Empty: pass + self._obs_queue.put((timestep, timed_obs.get_timestamp(), joint_state)) + + except Exception as e: + logger.exception(f"SendObservations processing error: {e}") + + return services_pb2.Empty() + + def GetActions(self, request, context): + try: + timestep, timestamp, joint_state = self._obs_queue.get(timeout=1.0) + except Empty: + return services_pb2.Actions(data=b"") + + if self._going_home: + # Episode just ended (max_episode_steps reached) โ€” drive to home instead + # of holding in place, for the duration of _auto_episode_end's pause. + action_deg = HOME_POSITION_DEG.copy() + else: + # Echo current joints back โ€” robot actively holds wherever it already is. + action_deg = joint_state.copy() + + timed_actions = [TimedAction( + timestamp=timestamp, timestep=timestep, + action=torch.tensor(action_deg, dtype=torch.float32), + )] + return services_pb2.Actions(data=pickle.dumps(timed_actions)) # nosec + + def _auto_episode_end(self): + """Auto-triggered once _episode_steps reaches max_episode_steps โ€” saves the + episode, drives the arm home (via GetActions checking _going_home), pauses + episode_pause_seconds, then resumes recording. Runs in a background thread so + the request/response loop keeps moving while the arm drives home and the + recorder flushes to disk. + """ + ep = self._episode_count + 1 + self._episode_count = ep + logger.info( + f"[ep {ep}] {self.max_episode_steps} steps reached โ€” saving and returning home " + f"(pausing {self.episode_pause_seconds:.1f}s before episode {ep + 1})..." + ) + self.recorder.flush() # slow disk write โ€” runs in background + self._going_home = True + time.sleep(self.episode_pause_seconds) + self._going_home = False + self._episode_steps = 0 + self.recorder.reset_pending() + self._recording_paused = False + logger.info(f"[ep {ep + 1}] recording started") + + def _episode_end(self): + """Triggered when the client calls Ready() with an episode already in + progress (i.e. ended early, before max_episode_steps) โ€” saves and waits for + an Enter keypress before resuming, instead of auto-resuming immediately. + """ + self._episode_count += 1 + # Pause recording immediately so observations sent while saving are discarded. + self._recording_paused = True + ep = self._episode_count + + def _save_and_wait(): + self.recorder.flush() # slow disk write โ€” runs in background + print( + f"\nโ”€โ”€ Episode {ep} saved โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n" + f" Press Enter to start recording episode {ep + 1} " + f"(Ctrl-C to stop)...", + flush=True, + ) + self._wait_for_enter() + + threading.Thread(target=_save_and_wait, daemon=True).start() + + def _wait_for_enter(self): + """Block in a background thread until the user presses Enter, then resume recording.""" + try: + input() + except EOFError: + pass + self.recorder.reset_pending() # clean slate for the new episode + self._recording_paused = False + print(f" Recording episode {self._episode_count + 1} ...", flush=True) + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# ENTRY POINT +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + +def serve(): + parser = argparse.ArgumentParser(description="Demo-recording gRPC server for SO-101") + parser.add_argument("--collect_dir", required=True, + help="Directory to save recorded (obs, next_obs) .npz files.") + parser.add_argument("--flush_every", type=int, default=200, + help="Flush buffer to disk every N steps (default: 200).") + parser.add_argument("--port", type=int, default=8080) + parser.add_argument("--max_episode_steps", type=int, default=200, + help="Hard cap on transitions recorded per episode โ€” auto-ends the " + "episode (save, return home, pause, resume) once reached.") + parser.add_argument("--episode_pause_seconds", type=float, default=5.0, + help="Pause this long (arm held at home) after each episode's " + "max_episode_steps is reached, before auto-resuming recording " + "for the next episode.") + parser.add_argument("--show_cameras", action="store_true", + help="Save live camera frames to cameras_latest.png during recording " + "(open in an editor that auto-refreshes to watch).") + args = parser.parse_args() + + recorder = DemoRecorder(Path(args.collect_dir), flush_every=args.flush_every) + logger.info(f"Demo recording enabled โ†’ {args.collect_dir} (flush every {args.flush_every} steps)") + + server = grpc.server(futures.ThreadPoolExecutor(max_workers=4)) + services_pb2_grpc.add_AsyncInferenceServicer_to_server( + DemoRecordingServicer( + recorder=recorder, + max_episode_steps=args.max_episode_steps, + show_cameras=args.show_cameras, + episode_pause_seconds=args.episode_pause_seconds, + ), + server, + ) + server.add_insecure_port(f"[::]:{args.port}") + server.start() + logger.info(f"gRPC server listening on port {args.port} mode=demo-recording") + logger.info("Waiting for robot_client.py ...") + try: + server.wait_for_termination() + except KeyboardInterrupt: + logger.info("Shutting down...") + recorder.flush() + logger.info("Final demo flush complete.") + server.stop(grace=2.0) + + +if __name__ == "__main__": + serve() diff --git a/mpail2/envs/real/so101/training/replay_demo.py b/mpail2/envs/real/so101/training/replay_demo.py new file mode 100644 index 0000000..9d10b54 --- /dev/null +++ b/mpail2/envs/real/so101/training/replay_demo.py @@ -0,0 +1,176 @@ +""" +Replay a recorded .npz demo on the physical SO-101 follower arm, +optionally showing the recorded camera images in a live window. + +Usage: + # Direct joint replay (original): + python replay_demo.py pick_changed/traj_0001.npz --port /dev/ttyACM0 --hz 10 + + # IK replay โ€” runs demo joints through FKโ†’IK to verify the Cartesian pipeline: + python replay_demo.py pick_changed/traj_0001.npz --port /dev/ttyACM0 --hz 10 --use_ik + + # Visual check only (no robot): + python replay_demo.py pick_changed/traj_0001.npz --no_robot --show_cameras + +--state_cols selects which columns of observation.state to use as joint targets. + Old 12-dim recordings: use 0 1 2 3 4 5 (follower, first 6). + New 6-dim recordings: use 0 1 2 3 4 5 (default, all 6). +""" + +import argparse +import time + +import numpy as np + + +JOINT_NAMES = [ + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", + "gripper", +] + + +def _to_uint8(img: np.ndarray) -> np.ndarray: + """Convert stored image (float32, any range) to uint8 RGB.""" + if img.dtype != np.uint8: + img = np.clip(img, 0, 255).astype(np.uint8) + return img + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("npz", help="Path to .npz demo file") + parser.add_argument("--port", default="/dev/ttyACM0", help="Follower arm serial port") + parser.add_argument("--robot_id", default="Kid") + parser.add_argument("--hz", type=float, default=10.0, help="Replay speed in Hz") + parser.add_argument( + "--state_cols", type=int, nargs="+", default=list(range(6)), + help="Which columns of observation.state to use (default: 0-5)" + ) + parser.add_argument( + "--show_cameras", action="store_true", + help="Display recorded camera images in a window during replay" + ) + parser.add_argument( + "--no_robot", action="store_true", + help="Skip robot connection โ€” only show camera images (useful for quick visual check)" + ) + parser.add_argument( + "--use_ik", action="store_true", + help="Replay via FKโ†’IK: convert demo joint positions to EE xyz then back to joints. " + "Verifies the Cartesian IK pipeline before live RL." + ) + args = parser.parse_args() + + data = np.load(args.npz, allow_pickle=False) + state = data["observation.state"] # (N, 2, D) + joint_positions = state[:, 0, args.state_cols] # (N, 6) โ€” obs side + n_steps = len(joint_positions) + + # Collect camera arrays present in the file + cam_keys = [k for k in data.keys() if k not in ("observation.state", "actions")] + cameras = {k: data[k][:, 0, :, :, :] for k in cam_keys} # (N, H, W, C) + + # Pre-compute IK targets if requested + ik_joint_positions = None + if args.use_ik: + import sys; sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent)) + from mpail2.envs.real.so101.ik_utils import fk, ik + print("Pre-computing IK targets from demo joint positions...") + ik_joints = [] + prev_arm = joint_positions[0, :5].copy() + for step_joints in joint_positions: + arm = step_joints[:5] + ee_xyz = fk(arm) + ik_arm = ik(ee_xyz, initial_arm_deg=prev_arm) + ik_joints.append(np.append(ik_arm, step_joints[5])) # keep original gripper + prev_arm = ik_arm + ik_joint_positions = np.array(ik_joints, dtype=np.float32) + + # Show round-trip error + ee_errors = [] + for orig, ik_j in zip(joint_positions, ik_joint_positions): + ee_orig = fk(orig[:5]) + ee_ik = fk(ik_j[:5]) + ee_errors.append(np.linalg.norm(ee_ik - ee_orig) * 1000) + print(f" IK EE round-trip error: mean={np.mean(ee_errors):.1f}mm max={np.max(ee_errors):.1f}mm") + + print(f"Loaded {args.npz}") + print(f" Steps: {n_steps}, mode={'IK' if args.use_ik else 'direct joints'}") + print(f" Cameras: {cam_keys if cam_keys else 'none'}") + print(f" Joint ranges:") + for i in range(len(args.state_cols)): + name = JOINT_NAMES[i] if i < len(JOINT_NAMES) else f"joint_{i}" + print(f" {name}: {joint_positions[:,i].min():.2f} to {joint_positions[:,i].max():.2f} deg") + + show = args.show_cameras and bool(cam_keys) + if args.show_cameras and not cam_keys: + print(" (no camera data in file โ€” --show_cameras ignored)") + + fig = ax = im_handle = None + if show: + import matplotlib + matplotlib.use("TkAgg") # use TkAgg; falls back gracefully if unavailable + import matplotlib.pyplot as plt + n_cams = len(cam_keys) + fig, axes = plt.subplots(1, n_cams, figsize=(6 * n_cams, 5)) + if n_cams == 1: + axes = [axes] + first_frames = [_to_uint8(cameras[k][0]) for k in cam_keys] + im_handles = [ax.imshow(f) for ax, f in zip(axes, first_frames)] + titles = [ax.set_title(k) for ax, k in zip(axes, cam_keys)] + step_text = fig.suptitle("step 0") + plt.tight_layout() + plt.pause(0.001) + + input(f"\nPress Enter to start replay at {args.hz} Hz ...") + + robot = None + if not args.no_robot: + from lerobot.robots.so_follower.config_so_follower import SOFollowerRobotConfig + from lerobot.robots import make_robot_from_config + cfg = SOFollowerRobotConfig(port=args.port, id=args.robot_id, use_degrees=True) + robot = make_robot_from_config(cfg) + robot.connect() + print("Robot connected. Starting replay...") + else: + print("--no_robot: skipping robot connection, showing images only.") + + replay_positions = ik_joint_positions if args.use_ik else joint_positions + + dt = 1.0 / args.hz + try: + for step, joints in enumerate(replay_positions): + t0 = time.perf_counter() + + if robot is not None: + action = {f"{name}.pos": float(joints[i]) for i, name in enumerate(JOINT_NAMES)} + robot.send_action(action) + + if show: + for im_h, k in zip(im_handles, cam_keys): + im_h.set_data(_to_uint8(cameras[k][step])) + step_text.set_text(f"step {step}/{n_steps} {'[IK]' if args.use_ik else ''}") + plt.pause(0.001) + + elapsed = time.perf_counter() - t0 + time.sleep(max(0.0, dt - elapsed)) + if step % 20 == 0: + print(f" step {step}/{n_steps} joints={joints.round(1)}") + + except KeyboardInterrupt: + print("\nReplay interrupted.") + finally: + if show: + import matplotlib.pyplot as plt + plt.close("all") + if robot is not None: + robot.disconnect() + print("Robot disconnected.") + + +if __name__ == "__main__": + main() diff --git a/mpail2/envs/real/so101/training/train_so101_local.py b/mpail2/envs/real/so101/training/train_so101_local.py new file mode 100644 index 0000000..978e4c0 --- /dev/null +++ b/mpail2/envs/real/so101/training/train_so101_local.py @@ -0,0 +1,273 @@ +""" +train_so101_local.py โ€” Train MPAIL2 on the real SO-101 via the local, non-dropping +gRPC env (mpail2/envs/real/so101), instead of the LeRobot async_inference / +demo_recording_server.py RPC-handler path. + +Unlike demo_recording_server.py (which reacts to inbound SendObservations/GetActions +RPCs from a separately-running lerobot robot_client), this script drives the +robot itself, in-process, via MPAIL2Runner.learn()'s standard rollout loop โ€” +env.step()/env.reset() block until the real action has been sent and the fresh +post-action state read back, so no observation is ever dropped. + +Requires mpail2.envs.real.so101.network.server to be running first (separate +`lerobot` conda env, connected to the real arm + cameras). See +mpail2/envs/real/so101/README.md for the full setup. + + conda activate lerobot + python -m mpail2.envs.real.so101.network.server --robot_port /dev/ttyACM0 --robot_id Kid \\ + --cam_index /dev/video0 --cam2_serial --grpc_port 7070 + +Then, in the `mpail2` conda env: + + conda activate mpail2 + python -m mpail2.envs.real.so101.training.train_so101_local --demo_path demo.pt --robot_host 127.0.0.1 \\ + --robot_port 7070 --device cuda --speed_scale 0.4 --wandb +""" + +import argparse +import logging +import os + +import numpy as np +import torch + +from mpail2.envs.real.so101 import ik_utils +from mpail2.envs.real.so101 import ( + SO101RealEnvArgs, make_so101_env, OBS_KEY, ACTION_DIM, EE_PROPRIO_DIM, +) +from mpail2.envs.real.so101.robot_limits import ( + CAM_KEY, CAM_H, CAM_W, CAM_C, CAM2_KEY, CAM2_H, CAM2_W, CAM2_C, +) +from mpail2.runner import MPAIL2Runner +from mpail2.configs.cfgs import MPAIL2RunnerCfg, ObsNormalizerCfg +from mpail2.configs.defs import ( + MultiCoderConfig, CNNCoderConfig, PlannerConfig, PolicySamplingConfig, LearnerConfig, +) + +# so101_env.py logs per-step action/observation details via logging.getLogger("so101_env"); +# this handler is what actually makes them visible on the console. +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + + +def main(): + parser = argparse.ArgumentParser( + description="Train MPAIL2 on the real SO-101 via the local gRPC env (no dropped observations)" + ) + parser.add_argument("--demo_path", required=True, help="Path to .pt demo file (same format as grpc_policy_server.py)") + parser.add_argument("--robot_host", default="127.0.0.1", help="so101_robot_server.py host") + parser.add_argument("--robot_port", type=int, default=7070, help="so101_robot_server.py gRPC port") + parser.add_argument("--device", default="cuda") + parser.add_argument("--log_dir", default="logs/so101_local") + parser.add_argument("--num_rollouts", type=int, default=512) + parser.add_argument("--num_elites", type=int, default=64) + parser.add_argument("--opt_iters", type=int, default=5) + parser.add_argument("--min_std", type=float, default=0.1, + help="MPPI sampling std floor. Doesn't help early on โ€” every reset/first " + "planning iter reinitializes _iter_std to max_std regardless.") + parser.add_argument("--max_std", type=float, default=3.0, + help="MPPI sampling std ceiling on the tanh-squashed [-1,1] action โ€” this is " + "the actual knob for 'how large is the raw action', independent of " + "speed_scale/lpf_alpha (which only rescale/smooth it afterward). An " + "untrained value fn won't shrink this much via CEM refinement, so a fresh " + "model will keep sampling near-saturated actions until it's had some " + "training. Lower this (e.g. 0.5-1.0) for gentler motion pre-training.") + parser.add_argument("--policy_proportion", type=float, default=0.05, + help="Fraction of MPPI rollouts drawn from the learned policy net vs pure " + "random exploration (default 0.05 = 95%% random, matches grpc_policy_server.py).") + parser.add_argument("--latent_dim", type=int, default=512) + parser.add_argument("--joint_dim", type=int, default=256) + parser.add_argument("--speed_scale", type=float, default=1.0) + parser.add_argument("--lpf_alpha", type=float, default=1.0, + help="Low-pass filter weight on the new action (0=frozen, 1=no filter). " + "Tames jerky/saturated MPPI output; speed_scale alone won't fix that " + "since it only rescales the envelope, not how often the action sits at it.") + parser.add_argument("--reward_scale", type=float, default=1.0) + parser.add_argument("--gripper_hold_steps", type=int, default=1, + help="Only accept a new gripper command every Nth step (holding the " + "current target fixed the rest of the time), so gripper motion is a " + "deliberate chunked open/close instead of raw per-step MPPI jitter. " + "1 = update every step (default, unchanged behavior).") + parser.add_argument("--max_episode_steps", type=int, default=200, + help="Steps per learn() iteration (one rollout before each update()).") + parser.add_argument("--reset_pause_seconds", type=float, default=3.0, + help="Pause after homing before the next episode's actions start, so the " + "arm actually settles instead of immediately being commanded again.") + parser.add_argument("--loss_horizon", type=int, default=7) + parser.add_argument("--replay_size", type=int, default=40_000) + parser.add_argument("--replay_batch_size", type=int, default=256) + parser.add_argument("--num_episodes", type=int, default=200, + help="Number of episodes to run. Each episode is one runner.learn() call " + "(max_episode_steps env steps + one update()), reset to home first.") + parser.add_argument("--checkpoint_every", type=int, default=20) + parser.add_argument("--load_checkpoint", default=None, help="Resume from a .pt checkpoint saved by MPAIL2Runner.save().") + parser.add_argument("--no_wrist_cam", action="store_true", help="Exclude wrist cam from the encoder.") + parser.add_argument("--wandb", action="store_true") + parser.add_argument("--wandb_project", default="so101-mpail2-local") + parser.add_argument("--wandb_run_name", default=None) + parser.add_argument("--eval", action="store_true", + help="Roll the loaded checkpoint out on the robot with no training updates, " + "no replay storage, no checkpoint saving โ€” just act() + env.step(). " + "Requires --load_checkpoint.") + parser.add_argument("--num_eval_episodes", type=int, default=10, + help="Number of episodes to roll out when --eval is set.") + parser.add_argument("--eval_policy_only", action="store_true", + help="During --eval, bypass CEM/MPPI entirely and use the policy network's " + "own deterministic mean action every step (Planner.act_policy_only) โ€” " + "100%% policy-driven, no noise rollouts, no elite re-weighting. Default " + "eval behavior otherwise still runs the full CEM search with " + "deterministic=True on the final iteration.") + args = parser.parse_args() + + if args.eval and not args.load_checkpoint: + parser.error("--eval requires --load_checkpoint.") + + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + device = ("cuda" if torch.cuda.is_available() else "cpu") if args.device == "auto" else args.device + + print(f"Loading demonstrations from {args.demo_path} ...") + demonstrations = torch.load(args.demo_path, map_location="cpu", weights_only=False) + demo_keys = {OBS_KEY, CAM_KEY, CAM2_KEY} + demonstrations = {k: v.float() for k, v in demonstrations.items() if k in demo_keys} + if CAM_KEY not in demonstrations: + raise RuntimeError(f"Demo file missing '{CAM_KEY}'. Keys: {list(demonstrations)}") + if CAM2_KEY not in demonstrations: + print(f"[WARN] Demo file missing '{CAM2_KEY}' โ€” training without RealSense camera.") + demonstrations[CAM2_KEY] = torch.zeros( + demonstrations[CAM_KEY].shape[:-1] + (CAM2_C,), dtype=torch.float32 + ) + + # No re-pairing/subsampling: at the current control rate + speed_scale, online + # per-step EE motion (MAX_DELTA_M * speed_scale) is already the same single-digit-mm + # order of magnitude as the demo's raw per-frame motion, so there's no longer a large + # scale gap for the discriminator to trivially exploit โ€” unlike when this was tuned + # against a 1Hz online control loop. Using demos in their raw recorded form also + # avoids re-pairing groups ever straddling two independently-recorded demo takes. + print(" Converting demo joint states to EE proprioception (running FK)...") + demo_joints = demonstrations[OBS_KEY].numpy() + N = demo_joints.shape[0] + ee_proprio = np.zeros((N, 2, EE_PROPRIO_DIM), dtype=np.float32) + for i in range(N): + for j in range(2): + ee_proprio[i, j] = ik_utils.joints_to_ee_proprio(demo_joints[i, j]) + demonstrations[OBS_KEY] = torch.from_numpy(ee_proprio) + for k, v in demonstrations.items(): + print(f" {k}: {tuple(v.shape)}") + + print(f"Connecting to so101_robot_server.py at {args.robot_host}:{args.robot_port} ...") + env = make_so101_env(SO101RealEnvArgs( + device=device, host=args.robot_host, port=args.robot_port, + max_episode_length=args.max_episode_steps, speed_scale=args.speed_scale, + lpf_alpha=args.lpf_alpha, reset_pause_seconds=args.reset_pause_seconds, + gripper_hold_steps=args.gripper_hold_steps, mock=False, + )) + + _coder_list = [ + MultiCoderConfig.ProprioCoderConfig(obs_key=OBS_KEY, input_dim=EE_PROPRIO_DIM, output_dim=args.joint_dim), + CNNCoderConfig(obs_key=CAM2_KEY, H=CAM2_H, W=CAM2_W, C=CAM2_C), + ] + if not args.no_wrist_cam: + _coder_list.append(CNNCoderConfig(obs_key=CAM_KEY, H=CAM_H, W=CAM_W, C=CAM_C)) + encoder_cfg = MultiCoderConfig(coder_list=_coder_list) + + planner_cfg = PlannerConfig( + encoder_cfg=encoder_cfg, action_dim=ACTION_DIM, latent_dim=args.latent_dim, + sampling_cfg=PolicySamplingConfig( + num_rollouts=args.num_rollouts, num_timesteps=args.loss_horizon, + policy_proportion=args.policy_proportion, min_std=args.min_std, max_std=args.max_std, + ), + opt_iters=args.opt_iters, num_elites=args.num_elites, + ) + learner_cfg = LearnerConfig( + planner_cfg=planner_cfg, replay_size=args.replay_size, replay_batch_size=args.replay_batch_size, + loss_horizon=args.loss_horizon, use_terminations=False, + obs_normalizer_cfg=ObsNormalizerCfg(normalization_type="cam_only"), + ) + + if args.wandb: + import wandb as _wandb + _wandb.init(project=args.wandb_project, name=args.wandb_run_name, config=vars(args)) + print(f"W&B run: {_wandb.run.url}") + + log_cfg = MPAIL2RunnerCfg.LogCfg( + log_dir=args.log_dir, checkpoint_every=args.checkpoint_every, + no_wandb=not args.wandb, logger="wandb" if args.wandb else None, video_interval=999_999, + ) + # num_learning_iterations is fixed at 1: runner.learn() calls env.reset() once at its + # own start (runner.py:108), then rolls out max_episode_steps steps and updates. We call + # learn() once per episode from an outer loop below, so each call's own reset() gives us + # per-episode homing "for free" without touching runner.py's shared learn() loop (which + # also backs Isaac-sim training, where individual sub-envs auto-reset internally โ€” a real + # robot doesn't, so calling learn() once for many iterations would never re-home the arm). + runner_cfg = MPAIL2RunnerCfg( + learner_cfg=learner_cfg, log_cfg=log_cfg, + num_learning_iterations=1, logger="wandb" if args.wandb else None, vis_rollouts=False, + ) + os.makedirs(os.path.join(args.log_dir, "models"), exist_ok=True) + runner = MPAIL2Runner(demonstrations=demonstrations, env=env, runner_cfg=runner_cfg, device=device) + + if args.load_checkpoint: + print(f"Loading checkpoint from {args.load_checkpoint} ...") + runner.load(args.load_checkpoint) + ckpt = torch.load(args.load_checkpoint, map_location="cpu", weights_only=False) + runner.current_learning_iteration = int(ckpt.get("iter", 0)) + print(f" Resuming from iteration {runner.current_learning_iteration}") + + runner.learner._reward_scale = float(args.reward_scale) + + # Per-step reward logging: env.step()'s own reward stays 0.0 (learner.py: "NOT TO BE USED + # IN LEARNING - JUST FOR LOGGING") โ€” the actual GAIL discriminator reward comparing + # consecutive latent embeddings is what grpc_policy_server.py computes inline for its + # [ep N step M] disc_reward= log lines. Wire the same computation into the wrapper here + # so it prints every step without touching runner.py's shared learn() loop. + def _disc_reward_fn(obs, next_obs): + with torch.no_grad(): + z = runner.learner._encoder(obs) + zn = runner.learner._encoder(next_obs) + return float(runner.learner._reward(z, zn).item()) + env.reward_fn = _disc_reward_fn + + print( + f"Runner ready device={device} proprio_dim={EE_PROPRIO_DIM} action_dim={ACTION_DIM} " + f"latent_dim={planner_cfg.latent_dim}" + ) + + if args.eval: + print(f"Evaluating checkpoint: {args.num_eval_episodes} episodes of {args.max_episode_steps} " + f"steps each, no training updates " + f"({'policy-net only, no CEM' if args.eval_policy_only else 'full CEM, deterministic final iter'}) ...") + runner.learner.eval() + for episode in range(args.num_eval_episodes): + print(f"[eval episode {episode + 1}/{args.num_eval_episodes}] resetting to home, then rolling out ...") + obs, info = env.reset() + runner.learner.planner.reset() + for _ in range(args.max_episode_steps): + with torch.no_grad(): + if args.eval_policy_only: + action = runner.learner.planner.act_policy_only(obs) + else: + action = runner.learner.planner.act(obs, deterministic=True) + obs, reward, terminated, truncated, info = env.step(action.to(env.unwrapped.device)) + if terminated or truncated: + break + return + + print(f"Starting training: {args.num_episodes} episodes of {args.max_episode_steps} steps each ...") + + start_episode = runner.current_learning_iteration + try: + for episode in range(start_episode, start_episode + args.num_episodes): + print(f"[episode {episode + 1}/{start_episode + args.num_episodes}] resetting to home, then rolling out ...") + runner.learn() + # runner.learn() always runs exactly one iteration (runner_cfg.num_learning_iterations=1 + # above) and its internal current_learning_iteration bookkeeping is a no-op for that + # case (it just resets to whatever it already was) โ€” advance it here instead, or every + # checkpoint keeps overwriting "model_1.pt" and wandb's it/tot_iter never move. + runner.current_learning_iteration = episode + 1 + except KeyboardInterrupt: + print("Interrupted โ€” saving checkpoint before exit...") + runner.save(postfix=f"ep{runner.current_learning_iteration}_interrupted") + raise + + +if __name__ == "__main__": + main() diff --git a/mpail2/envs/real/so101/transport/__init__.py b/mpail2/envs/real/so101/transport/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mpail2/envs/real/so101/transport/services_pb2.py b/mpail2/envs/real/so101/transport/services_pb2.py new file mode 100644 index 0000000..eba4057 --- /dev/null +++ b/mpail2/envs/real/so101/transport/services_pb2.py @@ -0,0 +1,49 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# Local copy of lerobot/transport/services_pb2.py for use in mpail2 env. +# source: lerobot/transport/services.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder + +try: + _runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, 31, 0, '', + 'lerobot/transport/services.proto' + ) +except Exception: + pass # allow mismatched protobuf versions + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n lerobot/transport/services.proto\x12\ttransport\"L\n\nTransition\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"L\n\nParameters\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"T\n\x12InteractionMessage\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"M\n\x0bObservation\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"\x17\n\x07\x41\x63tions\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x1b\n\x0bPolicySetup\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x07\n\x05\x45mpty*`\n\rTransferState\x12\x14\n\x10TRANSFER_UNKNOWN\x10\x00\x12\x12\n\x0eTRANSFER_BEGIN\x10\x01\x12\x13\n\x0fTRANSFER_MIDDLE\x10\x02\x12\x10\n\x0cTRANSFER_END\x10\x03\x32\x81\x02\n\x0eLearnerService\x12=\n\x10StreamParameters\x12\x10.transport.Empty\x1a\x15.transport.Parameters0\x01\x12<\n\x0fSendTransitions\x12\x15.transport.Transition\x1a\x10.transport.Empty(\x01\x12\x45\n\x10SendInteractions\x12\x1d.transport.InteractionMessage\x1a\x10.transport.Empty(\x01\x12+\n\x05Ready\x12\x10.transport.Empty\x1a\x10.transport.Empty2\xf5\x01\n\x0e\x41syncInference\x12>\n\x10SendObservations\x12\x16.transport.Observation\x1a\x10.transport.Empty(\x01\x12\x32\n\nGetActions\x12\x10.transport.Empty\x1a\x12.transport.Actions\x12\x42\n\x16SendPolicyInstructions\x12\x16.transport.PolicySetup\x1a\x10.transport.Empty\x12+\n\x05Ready\x12\x10.transport.Empty\x1a\x10.transport.Emptyb\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'lerobot.transport.services_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_TRANSFERSTATE']._serialized_start = 431 + _globals['_TRANSFERSTATE']._serialized_end = 527 + _globals['_TRANSITION']._serialized_start = 47 + _globals['_TRANSITION']._serialized_end = 123 + _globals['_PARAMETERS']._serialized_start = 125 + _globals['_PARAMETERS']._serialized_end = 201 + _globals['_INTERACTIONMESSAGE']._serialized_start = 203 + _globals['_INTERACTIONMESSAGE']._serialized_end = 287 + _globals['_OBSERVATION']._serialized_start = 289 + _globals['_OBSERVATION']._serialized_end = 366 + _globals['_ACTIONS']._serialized_start = 368 + _globals['_ACTIONS']._serialized_end = 391 + _globals['_POLICYSETUP']._serialized_start = 393 + _globals['_POLICYSETUP']._serialized_end = 420 + _globals['_EMPTY']._serialized_start = 422 + _globals['_EMPTY']._serialized_end = 429 + _globals['_LEARNERSERVICE']._serialized_start = 530 + _globals['_LEARNERSERVICE']._serialized_end = 787 + _globals['_ASYNCINFERENCE']._serialized_start = 790 + _globals['_ASYNCINFERENCE']._serialized_end = 1035 diff --git a/mpail2/envs/real/so101/transport/services_pb2_grpc.py b/mpail2/envs/real/so101/transport/services_pb2_grpc.py new file mode 100644 index 0000000..05e6e33 --- /dev/null +++ b/mpail2/envs/real/so101/transport/services_pb2_grpc.py @@ -0,0 +1,102 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +# Local copy for use in mpail2 env โ€” imports from local services_pb2 instead of lerobot. +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from mpail2.envs.real.so101.transport import services_pb2 as lerobot_dot_transport_dot_services__pb2 + +GRPC_GENERATED_VERSION = '1.73.1' +GRPC_VERSION = grpc.__version__ + +try: + from grpc._utilities import first_version_is_lower + if first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION): + warnings.warn( + f'grpc {GRPC_VERSION} is older than generated code target {GRPC_GENERATED_VERSION}. ' + 'Things may still work, but consider upgrading grpcio.', + RuntimeWarning, + ) +except ImportError: + pass + + +class AsyncInferenceStub(object): + """Client stub for AsyncInference service.""" + + def __init__(self, channel): + self.SendObservations = channel.stream_unary( + '/transport.AsyncInference/SendObservations', + request_serializer=lerobot_dot_transport_dot_services__pb2.Observation.SerializeToString, + response_deserializer=lerobot_dot_transport_dot_services__pb2.Empty.FromString, + ) + self.GetActions = channel.unary_unary( + '/transport.AsyncInference/GetActions', + request_serializer=lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString, + response_deserializer=lerobot_dot_transport_dot_services__pb2.Actions.FromString, + ) + self.SendPolicyInstructions = channel.unary_unary( + '/transport.AsyncInference/SendPolicyInstructions', + request_serializer=lerobot_dot_transport_dot_services__pb2.PolicySetup.SerializeToString, + response_deserializer=lerobot_dot_transport_dot_services__pb2.Empty.FromString, + ) + self.Ready = channel.unary_unary( + '/transport.AsyncInference/Ready', + request_serializer=lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString, + response_deserializer=lerobot_dot_transport_dot_services__pb2.Empty.FromString, + ) + + +class AsyncInferenceServicer: + """AsyncInference: from Robot perspective + Robot send observations to & executes action received from a remote Policy server + """ + + def SendObservations(self, request_iterator, context): + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetActions(self, request, context): + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def SendPolicyInstructions(self, request, context): + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Ready(self, request, context): + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_AsyncInferenceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'SendObservations': grpc.stream_unary_rpc_method_handler( + servicer.SendObservations, + request_deserializer=lerobot_dot_transport_dot_services__pb2.Observation.FromString, + response_serializer=lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString, + ), + 'GetActions': grpc.unary_unary_rpc_method_handler( + servicer.GetActions, + request_deserializer=lerobot_dot_transport_dot_services__pb2.Empty.FromString, + response_serializer=lerobot_dot_transport_dot_services__pb2.Actions.SerializeToString, + ), + 'SendPolicyInstructions': grpc.unary_unary_rpc_method_handler( + servicer.SendPolicyInstructions, + request_deserializer=lerobot_dot_transport_dot_services__pb2.PolicySetup.FromString, + response_serializer=lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString, + ), + 'Ready': grpc.unary_unary_rpc_method_handler( + servicer.Ready, + request_deserializer=lerobot_dot_transport_dot_services__pb2.Empty.FromString, + response_serializer=lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'transport.AsyncInference', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('transport.AsyncInference', rpc_method_handlers) diff --git a/mpail2/envs/real/so101/transport/so101_robot.proto b/mpail2/envs/real/so101/transport/so101_robot.proto new file mode 100644 index 0000000..d87cc02 --- /dev/null +++ b/mpail2/envs/real/so101/transport/so101_robot.proto @@ -0,0 +1,31 @@ +// gRPC service for the SO-101 real-robot control server (mpail2/envs/real/so101). +// +// Replaces the HTTP (/state, /step, /camera, /camera2, /reset) design described in +// so101_env.py's original docstring. One RPC per env.step()/env.reset() call instead +// of four separate HTTP round-trips (state + step + camera + camera2), and camera +// frames are sent as raw uint8 bytes instead of JSON-nested float lists. +// +// Regenerate after editing: +// python -m grpc_tools.protoc -I transport --python_out=transport --grpc_python_out=transport transport/so101_robot.proto + +syntax = "proto3"; +package so101robot; + +service SO101Robot { + // Move the arm to HOME_POSITION_DEG and open the gripper; returns the resulting state. + rpc Reset (ResetRequest) returns (RobotState); + // Command a joint-position target; returns the state resulting from that action. + rpc Step (StepRequest) returns (RobotState); +} + +message ResetRequest {} + +message StepRequest { + repeated float joints_deg = 1; // 6 target joint degrees, JOINT_NAMES order +} + +message RobotState { + repeated float joints_deg = 1; // 6 current joint degrees, JOINT_NAMES order + bytes cam = 2; // wrist cam, raw uint8, row-major (CAM_H, CAM_W, CAM_C) + bytes cam2 = 3; // RealSense cam, raw uint8, row-major (CAM2_H, CAM2_W, CAM2_C) +} diff --git a/mpail2/envs/real/so101/transport/so101_robot_pb2.py b/mpail2/envs/real/so101/transport/so101_robot_pb2.py new file mode 100644 index 0000000..f22c6f5 --- /dev/null +++ b/mpail2/envs/real/so101/transport/so101_robot_pb2.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: so101_robot.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 33, + 5, + '', + 'so101_robot.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x11so101_robot.proto\x12\nso101robot\"\x0e\n\x0cResetRequest\"!\n\x0bStepRequest\x12\x12\n\njoints_deg\x18\x01 \x03(\x02\";\n\nRobotState\x12\x12\n\njoints_deg\x18\x01 \x03(\x02\x12\x0b\n\x03\x63\x61m\x18\x02 \x01(\x0c\x12\x0c\n\x04\x63\x61m2\x18\x03 \x01(\x0c\x32\x80\x01\n\nSO101Robot\x12\x39\n\x05Reset\x12\x18.so101robot.ResetRequest\x1a\x16.so101robot.RobotState\x12\x37\n\x04Step\x12\x17.so101robot.StepRequest\x1a\x16.so101robot.RobotStateb\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'so101_robot_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_RESETREQUEST']._serialized_start=33 + _globals['_RESETREQUEST']._serialized_end=47 + _globals['_STEPREQUEST']._serialized_start=49 + _globals['_STEPREQUEST']._serialized_end=82 + _globals['_ROBOTSTATE']._serialized_start=84 + _globals['_ROBOTSTATE']._serialized_end=143 + _globals['_SO101ROBOT']._serialized_start=146 + _globals['_SO101ROBOT']._serialized_end=274 +# @@protoc_insertion_point(module_scope) diff --git a/mpail2/envs/real/so101/transport/so101_robot_pb2_grpc.py b/mpail2/envs/real/so101/transport/so101_robot_pb2_grpc.py new file mode 100644 index 0000000..57a0043 --- /dev/null +++ b/mpail2/envs/real/so101/transport/so101_robot_pb2_grpc.py @@ -0,0 +1,144 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +# Local copy for use in both mpail2 and lerobot envs โ€” imports from local +# so101_robot_pb2 instead of a bare top-level module. +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from mpail2.envs.real.so101.transport import so101_robot_pb2 as so101__robot__pb2 + +GRPC_GENERATED_VERSION = '1.81.1' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in so101_robot_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class SO101RobotStub: + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Reset = channel.unary_unary( + '/so101robot.SO101Robot/Reset', + request_serializer=so101__robot__pb2.ResetRequest.SerializeToString, + response_deserializer=so101__robot__pb2.RobotState.FromString, + _registered_method=True) + self.Step = channel.unary_unary( + '/so101robot.SO101Robot/Step', + request_serializer=so101__robot__pb2.StepRequest.SerializeToString, + response_deserializer=so101__robot__pb2.RobotState.FromString, + _registered_method=True) + + +class SO101RobotServicer: + """Missing associated documentation comment in .proto file.""" + + def Reset(self, request, context): + """Move the arm to HOME_POSITION_DEG and open the gripper; returns the resulting state. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Step(self, request, context): + """Command a joint-position target; returns the state resulting from that action. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_SO101RobotServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Reset': grpc.unary_unary_rpc_method_handler( + servicer.Reset, + request_deserializer=so101__robot__pb2.ResetRequest.FromString, + response_serializer=so101__robot__pb2.RobotState.SerializeToString, + ), + 'Step': grpc.unary_unary_rpc_method_handler( + servicer.Step, + request_deserializer=so101__robot__pb2.StepRequest.FromString, + response_serializer=so101__robot__pb2.RobotState.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'so101robot.SO101Robot', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('so101robot.SO101Robot', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class SO101Robot: + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def Reset(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/so101robot.SO101Robot/Reset', + so101__robot__pb2.ResetRequest.SerializeToString, + so101__robot__pb2.RobotState.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def Step(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/so101robot.SO101Robot/Step', + so101__robot__pb2.StepRequest.SerializeToString, + so101__robot__pb2.RobotState.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/mpail2/envs/real/so101/wrappers.py b/mpail2/envs/real/so101/wrappers.py new file mode 100644 index 0000000..c0c3bf7 --- /dev/null +++ b/mpail2/envs/real/so101/wrappers.py @@ -0,0 +1,111 @@ +"""MPAIL-shaped wrapper for SO101RobotEnv (mirrors FrankaRealWrapper / KinovaRealWrapper).""" + +from __future__ import annotations + +import logging +import time +from typing import Any, Callable, Dict, Optional, Tuple + +import gymnasium as gym +import numpy as np +import torch + +from .robot_limits import MAX_EPISODE_STEPS + +logger = logging.getLogger("so101_env") + + +class SO101RealWrapper(gym.Wrapper): + """Converts SO101RobotEnv numpy observations to torch tensors. + + Sets ``num_envs`` and ``max_episode_length`` on the *inner* env so that + ``env.unwrapped`` exposes both attributes as required by MPAIL2Runner. + """ + + def __init__(self, env: gym.Env, device: str = "cuda"): + super().__init__(env) + self.device = device + self.step_count = 0 + self.episode_count = 0 + + # Optional callback(prev_obs_torch, obs_torch) -> float, set by the caller after + # building the learner (e.g. train_so101_local.py) to compute + log the actual + # GAIL discriminator reward per step, the same way grpc_policy_server.py does + # inline via learner._encoder/_reward. runner.learn() itself never sees this โ€” + # the env's own step() reward stays 0.0, since that's purely a logging value + # (learner.py: "THESE ARE NOT TO BE USED IN LEARNING - JUST FOR LOGGING"). + self.reward_fn: Optional[Callable[[Dict[str, torch.Tensor], Dict[str, torch.Tensor]], float]] = None + self._prev_obs_torch: Optional[Dict[str, torch.Tensor]] = None + + # Expose these on unwrapped for MPAIL2Runner + self.env.unwrapped.num_envs = 1 + self.env.unwrapped.device = device + if not hasattr(self.env.unwrapped, "max_episode_length"): + self.env.unwrapped.max_episode_length = MAX_EPISODE_STEPS + self.max_episode_length = self.env.unwrapped.max_episode_length + + # โ”€โ”€โ”€ observation helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def _obs_to_torch( + self, obs: Dict[str, np.ndarray] + ) -> Dict[str, torch.Tensor]: + return { + k: torch.as_tensor(v, dtype=torch.float32, device=self.device) + for k, v in obs.items() + } + + # โ”€โ”€โ”€ Gymnasium API โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def reset( + self, + *, + seed: int | None = None, + options: Dict[str, Any] | None = None, + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: + self.step_count = 0 + self.episode_count += 1 + self._prev_obs_torch = None # discriminator reward needs a real (obs_t, obs_t+1) pair + t0 = time.time() + obs, info = self.env.reset(seed=seed, options=options) + info = dict(info) + info["reset_time"] = info.get("reset_time", round(time.time() - t0, 3)) + obs_torch = self._obs_to_torch(obs) + self._prev_obs_torch = obs_torch + return obs_torch, info + + def step( + self, action: torch.Tensor | np.ndarray + ) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, Any]]: + t0 = time.time() + + if isinstance(action, torch.Tensor): + action_np = action.detach().cpu().numpy() + else: + action_np = np.asarray(action, dtype=np.float32) + + obs, reward, terminated, truncated, info = self.env.step(action_np) + obs_torch = self._obs_to_torch(obs) + + self.step_count += 1 + truncated = truncated or (self.step_count >= self.max_episode_length) + + info = dict(info) + info["action_executed"] = action_np.flatten().tolist() + info["mpail_env/step_time"] = round(time.time() - t0, 3) + + if self.reward_fn is not None and self._prev_obs_torch is not None: + disc_reward = self.reward_fn(self._prev_obs_torch, obs_torch) + info["disc_reward"] = disc_reward + logger.info(f"[step {self.step_count - 1}] disc_reward={disc_reward:.4f}") + self._prev_obs_torch = obs_torch + + return ( + obs_torch, + torch.tensor([float(reward)], dtype=torch.float32, device=self.device), + torch.tensor([bool(terminated)], dtype=torch.bool, device=self.device), + torch.tensor([bool(truncated)], dtype=torch.bool, device=self.device), + info, + ) + + def seed(self, seed: int | None = None) -> None: + pass diff --git a/mpail2/learner.py b/mpail2/learner.py index 5e228b6..2fcec26 100644 --- a/mpail2/learner.py +++ b/mpail2/learner.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Dict, Optional from mpail2.encoder import Coder -from mpail2.dynamics import Dynamics +from mpail2.dynamics import Dynamics, SIGReg from mpail2.reward import Reward from mpail2.value import EnsembleValue from mpail2.sampling import PolicyNetwork @@ -111,6 +111,14 @@ def __init__(self, **self.dynamics_learner_cfg.opt_params ) + self._sigreg = None + if getattr(self.dynamics_learner_cfg, 'sigreg_coeff', None) is not None \ + and self.dynamics_learner_cfg.sigreg_coeff > 0.0: + self._sigreg = SIGReg( + knots=self.dynamics_learner_cfg.sigreg_knots, + num_proj=self.dynamics_learner_cfg.sigreg_num_proj, + ).to(device=self.device, dtype=self.dtype) + # # POLICY SAMPLING SETUP # @@ -400,6 +408,16 @@ def update_dynamics( _jep_se = (pred_latents[:, 1:, :] - next_latent_batch_traj.detach()).pow(2) loss = (_rhos[..., None] * _jep_se).mean() + sigreg_loss = torch.tensor(0.0, device=self.device, dtype=self.dtype) + if self._sigreg is not None: + # SIGReg expects (T, B, D); latent_batch_traj (the gradient-carrying encode + # of the full obs_batch_traj horizon, not just z0) already has this shape + # once transposed. + sigreg_proj = latent_batch_traj.transpose(0, 1) + sigreg_loss = self._sigreg(sigreg_proj) + loss = loss + self.dynamics_learner_cfg.sigreg_coeff * sigreg_loss + self._mean_stats['Dyn/sigreg_loss'] += sigreg_loss.item() + # Decoder for visualization if specified if self._decoder: _recon_loss = torch.tensor(0.0, device=self.device, dtype=self.dtype) @@ -552,7 +570,8 @@ def update_value( # TODO: share policy-planned trajectories with policy update """ with torch.no_grad(): - rewards_batch = self._reward(latent_batch, next_latent_batch, action=action_batch) + _rscale = getattr(self, '_reward_scale', 1.0) + rewards_batch = self._reward(latent_batch, next_latent_batch, action=action_batch) * _rscale _next_plans = self._policy.plan(next_latent_batch) # [batch_size, H, action_dim] pred_next_latent_batch_traj = self._dynamics( z0=next_latent_batch, controls=_next_plans, @@ -665,9 +684,10 @@ def _n_step_return_lambda( ''' _gam, _H = self.value_learner_cfg.gamma, self.cfg.loss_horizon + _rscale = getattr(self, '_reward_scale', 1.0) _all_rewards = reward_fn( # [batch, H] z_traj, next_z_traj, action=actions_traj - ) + ) * _rscale _all_vals = value(z_traj, actions_traj, return_type=value_return_type) # [batch, H] or [num_q, batch, H] if log_probs is not None: # Entropy regularization adjustment diff --git a/mpail2/planner.py b/mpail2/planner.py index 27343cc..37f41e8 100644 --- a/mpail2/planner.py +++ b/mpail2/planner.py @@ -205,6 +205,27 @@ def act( return actions + def act_policy_only(self, observations: Dict[str, torch.Tensor]) -> torch.Tensor: + '''Return the policy network's own deterministic (tanh-squashed mean) action, + bypassing CEM/MPPI entirely โ€” no noise rollouts, no elite re-weighting, no + blending with the exploration distribution. For evaluating what the learned + policy net alone has learned, independent of the planner's search process. + ''' + if self._obs_normalizer is not None: + observations = self._obs_normalizer(observations) + + with torch.no_grad(): + z0 = self.encoder(observations) + self.sampling.policy.plan(z0) + actions = self.sampling.policy.action_mean + + self._last_actions = actions + self._current_obs = observations + with torch.no_grad(): + self._prev_z = z0 + + return actions + def update(self, obs: torch.Tensor, map: torch.Tensor=None): ''' Update the internal belief state of the agent. @@ -241,7 +262,15 @@ def step(self, obs, use_prev_opt:bool=True, deterministic:bool=False) -> torch.T else: self._opt_controls[:] = 0. - self.sampling.reset_iter_state() + # Bug fix: this used to call reset_iter_state() with no arguments, which resets + # _iter_mean to zero every single real-world step โ€” discarding the warm-start + # above and making ~95% of each decision's candidates (the noise-sampled ones, + # not the small policy_proportion fraction) a fresh zero-mean/max-std blind + # search with no memory of the previous decision's converged plan. Passing + # _opt_controls through means the noise samples are actually centered on the + # carried-forward plan, as intended (matches this method's own docstring: "resets + # _iter_std to max_std and _iter_mean to prev_controls"). + self.sampling.reset_iter_state(prev_controls=self._opt_controls) for i in range(self.cfg.opt_iters - 1): # Subsequent optimization uses previous optimal controls diff --git a/mpail2/reward.py b/mpail2/reward.py index d765d27..77f7122 100644 --- a/mpail2/reward.py +++ b/mpail2/reward.py @@ -37,4 +37,7 @@ def __init__( def forward(self, state, next_state, action=None): '''state shape: (num_envs, state_dim)''' _input = torch.cat([state, next_state], dim=-1) - return self.model(_input).squeeze(-1) + out = self.model(_input).squeeze(-1) + if self.cfg.reward_clip is not None: + out = out.clamp(-self.cfg.reward_clip, self.cfg.reward_clip) + return out diff --git a/mpail2/utils/obs_normalizer.py b/mpail2/utils/obs_normalizer.py index 30e4303..653a83a 100644 --- a/mpail2/utils/obs_normalizer.py +++ b/mpail2/utils/obs_normalizer.py @@ -34,6 +34,11 @@ def __init__( # Compute and register statistics from demonstrations self._compute_statistics(demonstrations) + @staticmethod + def _buf_name(key: str) -> str: + """Sanitize observation key to a valid buffer name (no dots allowed).""" + return key.replace(".", "_") + def _compute_statistics(self, demonstrations: Dict[str, torch.Tensor]): """Compute mean and std from expert demonstrations.""" self._obs_keys = list(demonstrations.keys()) @@ -52,9 +57,10 @@ def _compute_statistics(self, demonstrations: Dict[str, torch.Tensor]): mean = flat_data.mean(dim=0) std = flat_data.std(dim=0) + buf = self._buf_name(key) # Register as buffers so they're saved/loaded with model and moved with .to() - self.register_buffer(f"{key}_mean", mean.to(self.device)) - self.register_buffer(f"{key}_std", std.to(self.device)) + self.register_buffer(f"{buf}_mean", mean.to(self.device)) + self.register_buffer(f"{buf}_std", std.to(self.device)) print(f"[FixedObsNormalizer] {key}: mean range [{mean.min():.3f}, {mean.max():.3f}], " f"std range [{std.min():.3f}, {std.max():.3f}]") @@ -63,9 +69,10 @@ def forward(self, obs: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Union[to if isinstance(obs, dict): normalized_obs = {} for key, value in obs.items(): - if hasattr(self, f"{key}_mean"): - mean = getattr(self, f"{key}_mean") - std = getattr(self, f"{key}_std") + buf = self._buf_name(key) + if hasattr(self, f"{buf}_mean"): + mean = getattr(self, f"{buf}_mean") + std = getattr(self, f"{buf}_std") normalized_obs[key] = (value - mean) / (std + self.eps) if self.clip_obs is not None: normalized_obs[key] = torch.clamp(normalized_obs[key], min=-self.clip_obs, max=self.clip_obs) @@ -76,9 +83,9 @@ def forward(self, obs: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Union[to else: # For single tensor, use first key's statistics (assumes single obs type) if len(self._obs_keys) > 0: - key = self._obs_keys[0] - mean = getattr(self, f"{key}_mean") - std = getattr(self, f"{key}_std") + buf = self._buf_name(self._obs_keys[0]) + mean = getattr(self, f"{buf}_mean") + std = getattr(self, f"{buf}_std") return (obs - mean) / (std + self.eps) return obs @@ -86,18 +93,19 @@ def inverse(self, normalized_obs: Union[torch.Tensor, Dict[str, torch.Tensor]]) if isinstance(normalized_obs, dict): denormalized_obs = {} for key, value in normalized_obs.items(): - if hasattr(self, f"{key}_mean"): - mean = getattr(self, f"{key}_mean") - std = getattr(self, f"{key}_std") + buf = self._buf_name(key) + if hasattr(self, f"{buf}_mean"): + mean = getattr(self, f"{buf}_mean") + std = getattr(self, f"{buf}_std") denormalized_obs[key] = value * (std + self.eps) + mean else: denormalized_obs[key] = value return denormalized_obs else: if len(self._obs_keys) > 0: - key = self._obs_keys[0] - mean = getattr(self, f"{key}_mean") - std = getattr(self, f"{key}_std") + buf = self._buf_name(self._obs_keys[0]) + mean = getattr(self, f"{buf}_mean") + std = getattr(self, f"{buf}_std") return normalized_obs * (std + self.eps) + mean return normalized_obs @@ -115,8 +123,14 @@ def forward(self, obs: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Union[to normalized_obs = {} for key, value in obs.items(): if isinstance(value, torch.Tensor) and "cam" in key.lower(): - # Camera observations: normalize to [-0.5, 0.5] - normalized_obs[key] = value / 255.0 - 0.5 + # Camera observations: normalize to [-0.5, 0.5]. Some callers (e.g. + # so101_env.py, convert.py) already divide by 255 before this point, + # so dividing unconditionally here would silently double-divide and + # crush the whole batch into a ~[0, 1/255] sliver near -0.5 (near-zero + # variance regardless of image content). Only divide if still raw uint8 + # scale, mirroring the same max()>1.0 check convert.py already uses. + scaled = value / 255.0 if value.max() > 1.0 else value + normalized_obs[key] = scaled - 0.5 else: # Pass through all other observations unchanged normalized_obs[key] = value @@ -131,8 +145,11 @@ def inverse(self, normalized_obs: Union[torch.Tensor, Dict[str, torch.Tensor]]) denormalized_obs = {} for key, value in normalized_obs.items(): if isinstance(value, torch.Tensor) and "cam" in key.lower(): - # Camera observations: denormalize from [-0.5, 0.5] - denormalized_obs[key] = (value + 0.5) * 255.0 + # Undo forward()'s "- 0.5" back to [0, 1] โ€” matches this codebase's + # actual convention (so101_env.py/convert.py always hand forward() + # already-[0,1] images, so forward() never multiplies by 255 in + # practice; scaling back up to [0,255] here would be the wrong inverse). + denormalized_obs[key] = value + 0.5 else: # Pass through all other observations unchanged denormalized_obs[key] = value diff --git a/pyproject.toml b/pyproject.toml index 1db34ee..b23a14e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,13 @@ franka = [ "franky-control", "pyrealsense2", ] +so101 = [ + "opencv-python", + "pyrealsense2", + "grpcio", + "grpcio-tools", + "ikpy", +] [tool.setuptools.packages.find] include = ["mpail2*"]