vlastudio 0.2.0__py3-none-any.whl

vlastudio/__init__.py

@@ -0,0 +1,33 @@
+"""VLAStudio's lightweight public API. Importing it never installs dependencies."""
+from .extensions import register, resolve, create
+from .dependencies import DependencyError, environment_plan
+from .api import (load_dataset, load_policy, load_env, connect_policy, train, serve,
+                  show_camera, collect, infer,
+                  Dataset, Policy, RemotePolicy, Environment, TrainingResult,
+                  EvaluationResult, TaskError)
+
+__version__ = "0.2.0"
+__all__ = ["register", "resolve", "create", "load_dataset", "load_policy", "load_env",
+           "connect_policy", "train", "serve", "show_camera", "collect", "infer",
+           "Dataset", "Policy", "RemotePolicy", "Environment", "TrainingResult",
+           "EvaluationResult", "TaskError"]
+
+
+def run(command, args=(), **options):
+    """Run a task in the calling Python environment; return its exit code.
+
+    Example: run("train", ["-p", "my_policy.yaml", "-o", "/checkpoints/run"])
+    """
+    from .cli import main
+    argv = [command, *map(str, args)]
+    for key, value in options.items():
+        flag = "--" + key.replace("_", "-")
+        if value is True:
+            argv.append(flag)
+        elif value is not False and value is not None:
+            values = value if isinstance(value, (tuple, list)) else [value]
+            for item in values:
+                argv.extend([flag, str(item)])
+    return main(argv)
+
+__all__.extend(["run", "DependencyError", "environment_plan"])

vlastudio/__main__.py

@@ -0,0 +1,2 @@
+from .cli import main
+raise SystemExit(main())

vlastudio/api.py

@@ -0,0 +1,369 @@
+"""Configuration handles for training and evaluation workers.
+
+These are deliberately not torch modules or iterable datasets. They can be
+created in a lightweight Python process without importing policy dependencies.
+"""
+from dataclasses import dataclass, field
+import json
+from pathlib import Path
+from typing import Mapping
+
+from .configuration import resolve_config
+
+
+def _path(value):
+    return Path(value).expanduser().resolve()
+
+
+def _config(value, category):
+    return resolve_config(str(value), category)
+
+
+def _overrides(values):
+    result = []
+    for key, value in (values or {}).items():
+        if not isinstance(key, str) or not key or key.startswith('-'):
+            raise ValueError("Override keys must be argument names without leading dashes")
+        if isinstance(value, (dict, list, tuple)) or value is None:
+            raise TypeError("Overrides must be scalar; put structured values in YAML")
+        result.extend(["--" + key, str(value).lower() if isinstance(value, bool) else str(value)])
+    return result
+
+
+class TaskError(RuntimeError):
+    """A worker failed; logs are streamed to the calling process's terminal."""
+
+
+def _run(command, args, options):
+    from . import run
+    code = run(command, args, **options)
+    if code:
+        raise TaskError(f"{command} failed with exit code {code}; see worker output above")
+
+
+@dataclass(frozen=True)
+class Dataset:
+    """A task configuration, including dataset mixtures and normalization metadata."""
+    config_path: Path
+    cache_dir: Path | None = None
+
+
+@dataclass
+class Policy:
+    """A policy configuration and, after training, its saved checkpoint location."""
+    config_path: Path
+    checkpoint: Path | None = None
+    runtime_options: dict = field(default_factory=dict, repr=False)
+
+
+@dataclass(frozen=True)
+class RemotePolicy:
+    """A policy server address consumed by an evaluation worker."""
+    address: str
+
+
+@dataclass(frozen=True)
+class TrainingResult:
+    checkpoint: Path
+    policy: Policy
+
+
+@dataclass(frozen=True)
+class EvaluationResult:
+    output_dir: Path
+    metrics: dict
+
+
+@dataclass
+class Environment:
+    """Simulation evaluation configuration; construction does not start a simulator."""
+    config_path: Path
+    runtime_options: dict = field(default_factory=dict, repr=False)
+
+    def evaluate(self, policy, *, output_dir, num_rollout=4, batch_size=0,
+                 device="cuda", action_manager=None, overrides=None, **runtime_options):
+        """Evaluate a saved policy using the simulation evaluation pipeline.
+
+        Built-in environments declare their simulator dependencies. A custom
+        environment may supply runtime_manifest or use runtime='current'.
+        """
+        if isinstance(policy, str):
+            policy = connect_policy(policy)
+        if isinstance(policy, RemotePolicy):
+            model = policy.address
+            policy_options = {}
+        elif isinstance(policy, Policy) and policy.checkpoint is not None:
+            model = str(policy.checkpoint)
+            policy_options = {k: v for k, v in policy.runtime_options.items()
+                              if k != 'runtime_manifest'}
+        else:
+            raise ValueError("Evaluation requires a Policy checkpoint, RemotePolicy, or server address")
+        if num_rollout <= 0 or batch_size < 0:
+            raise ValueError("num_rollout must be positive and batch_size nonnegative")
+        reserved = {'o', 'output_dir', 'm', 'model_name_or_path', 'e', 'env',
+                    'n', 'num_rollout', 'bs', 'batch_size', 'device', 'am', 'action_manager'}
+        if reserved.intersection(overrides or {}):
+            raise ValueError("Use evaluate() parameters for output, checkpoint, environment and rollout options")
+        options = {**policy_options, **self.runtime_options,
+                   **_normalize_options(runtime_options)}
+        output = _path(output_dir)
+        if output.exists() and any(output.iterdir()):
+            raise ValueError("Use an empty evaluation output directory to avoid mixing old metrics")
+        args = ["-m", model, "-e", str(self.config_path),
+                "-o", str(output), "-n", str(num_rollout), "-bs", str(batch_size), "--device", device]
+        if action_manager is not None:
+            args.extend(["-am", str(_config(action_manager, "action_manager"))])
+        args.extend(_overrides(overrides))
+        _run("evaluate", args, options)
+        metrics = {str(p.relative_to(output)): json.loads(p.read_text(encoding="utf-8"))
+                   for p in sorted(output.rglob("*.json"))}
+        return EvaluationResult(output, metrics)
+
+
+def load_dataset(config="sim_transfer_cube_scripted", *, cache_dir=None, check_dependencies=True):
+    """Select a task alias or custom YAML, with an optional dataset-only cache.
+
+    Examples: 'sim_transfer_cube_scripted', 'rlbench.reach_target', '/my/task.yaml'.
+    None defers to the runtime's default data cache; no directories are created here.
+    Missing requirements raise DependencyError with install commands. Use
+    check_dependencies=False for configuration-only orchestration.
+    """
+    path = _config(config, "task")
+    if check_dependencies:
+        from .dependencies import require_dependencies
+        require_dependencies(dataset=str(path))
+    return Dataset(path, _path(cache_dir) if cache_dir is not None else None)
+
+
+def load_policy(config="act", *, checkpoint=None, check_dependencies=True, **runtime_options):
+    """Describe a policy using the same selector as train.py --policy.
+
+    checkpoint is for evaluation. To initialize training weights use the policy
+    config's pretrained_weight_path; to resume use training overrides.
+    Current-runtime calls check declared package metadata without importing the
+    model. Installation is always performed externally with pip or uv.
+    """
+    options = _normalize_options(runtime_options)
+    options.setdefault("runtime", "current")
+    path = _config(config, "policy")
+    if check_dependencies and options['runtime'] == 'current':
+        from .dependencies import require_dependencies
+        require_dependencies(**({'manifest': options['runtime_manifest']} if options.get('runtime_manifest') else {'policy': str(path)}))
+    return Policy(path, _path(checkpoint) if checkpoint else None, options)
+
+
+def connect_policy(address):
+    """Create a remote policy handle for TCP, HTTP(S), or shared memory."""
+    from .deploy.comm import is_server_address
+    value = str(address)
+    if not is_server_address(value):
+        raise ValueError(
+            "Policy server address must be host:port, http(s)://host:port, or shm://name"
+        )
+    return RemotePolicy(value)
+
+
+def load_env(config, *, check_dependencies=True, **runtime_options):
+    """Describe a benchmark and explain missing dependencies in the current runtime."""
+    options = _normalize_options(runtime_options)
+    options.setdefault("runtime", "current")
+    path = _config(config, "env")
+    if check_dependencies and options['runtime'] == 'current':
+        from .dependencies import require_dependencies
+        require_dependencies(**({'manifest': options['runtime_manifest']} if options.get('runtime_manifest') else {'environment': str(path)}))
+    return Environment(path, options)
+
+
+def show_camera(config, *, camera_names=None, fps=30.0, title="Camera",
+                window_scale=1.0, max_frames=None, show=True,
+                enable_viewer=False, **runtime_options):
+    """Start a (simulated) robot device and display its live camera feed.
+
+    The robot device runs in a subprocess, publishes observations to its shared
+    memory, and this function shows the camera frames in a window in real time.
+
+    Platform / headless behaviour (see ``vlastudio.deploy.camera_viewer``):
+    - Windows opens a window.
+    - POSIX opens a window only when ``DISPLAY`` is set (physical X server or
+      Xvfb); a headless session without ``DISPLAY`` reads frames without a
+      window. Pass ``show=False`` to force headless.
+
+    Args:
+        config: Robot config name or YAML path (same as ``vlastudio collect -r``).
+        camera_names: Which camera keys to display; default all present per frame.
+        fps: Target display frame rate.
+        title: OpenCV window title.
+        window_scale: Scale factor for the displayed frames.
+        max_frames: Stop after this many frames (None = run until interrupted).
+        show: Set False to never open a window.
+        enable_viewer: Open the MuJoCo 3D viewer as well (default off).
+
+    Returns:
+        Number of frames processed.
+    """
+    import multiprocessing as mp
+
+    from .configs.loader import ConfigLoader
+    from .deploy.base import _get_device_shm_name, start_device
+    from .deploy.camera_viewer import LiveCameraViewer
+
+    cfg, path = ConfigLoader().load_robot(config)
+    if isinstance(cfg, list):
+        # Multi-row YAML: pick the robot device, skip visualizer rows.
+        from .deploy.utils import is_visualizer_config
+
+        robot_rows = [row for row in cfg if isinstance(row, dict) and not is_visualizer_config(row)]
+        if not robot_rows:
+            raise ValueError(f"No robot device found in config: {path}")
+        cfg = robot_rows[0]
+    cfg = dict(cfg)
+    args = dict(cfg.get("args", {}))
+    args["enable_viewer"] = bool(enable_viewer)  # camera feed only unless asked
+    cfg["args"] = args
+    shm_name = _get_device_shm_name(cfg)
+
+    print(f"[show_camera] robot config: {path}", flush=True)
+    print(f"[show_camera] reading observations from SHM '{shm_name}'", flush=True)
+
+    proc = mp.Process(target=start_device, args=(cfg,), daemon=True)
+    proc.start()
+    try:
+        viewer = LiveCameraViewer(
+            shm_name,
+            camera_names=camera_names,
+            fps=fps,
+            title=title,
+            window_scale=window_scale,
+            show=show,
+        )
+        return viewer.run(max_frames=max_frames)
+    finally:
+        proc.terminate()
+        proc.join(timeout=2.0)
+        if proc.is_alive():
+            proc.kill()
+
+
+def serve(policy, *, address="0.0.0.0", port=None, device="cuda",
+          dataset_id="", chunk_size=-1, **runtime_options):
+    """Serve a local policy checkpoint until interrupted.
+
+    address accepts a TCP bind host, host:port, http(s)://host:port, or shm://name.
+    """
+    if not isinstance(policy, Policy) or policy.checkpoint is None:
+        raise ValueError("serve() requires a Policy with a supplied or trained checkpoint")
+    from .deploy.comm import is_http_address, is_server_address, is_shm_address, parse_server_address
+    bind = str(address)
+    if is_server_address(bind) and not is_http_address(bind) and not is_shm_address(bind):
+        bind, address_port = parse_server_address(bind)
+        if port is not None and port != address_port:
+            raise ValueError("port conflicts with the port included in address")
+        port = address_port
+    args = ["-m", str(policy.checkpoint), "--host", bind, "--device", str(device),
+            "--dataset_id", str(dataset_id), "--chunk_size", str(chunk_size)]
+    if port is not None:
+        if not 1 <= int(port) <= 65535:
+            raise ValueError("port must be between 1 and 65535")
+        args.extend(["--port", str(port)])
+    options = {**policy.runtime_options, **_normalize_options(runtime_options)}
+    _run("serve", args, options)
+
+
+def _normalize_options(options):
+    if options.get("runtime", "current") != "current":
+        raise ValueError("Only the current Python environment is supported. Install dependencies externally with pip or uv.")
+    allowed = {"runtime", "runtime_manifest", "cache_dir", "model_cache_dir",
+               "plugin_path", "config_path", "offline", "data_cache_dir"}
+    unknown = options.keys() - allowed
+    if unknown:
+        raise TypeError(f"Unknown runtime options: {sorted(unknown)}")
+    result = dict(options)
+    for key in ("runtime_manifest", "cache_dir", "model_cache_dir", "data_cache_dir"):
+        if result.get(key) is not None:
+            result[key] = str(_path(result[key]))
+    for key in ("plugin_path", "config_path"):
+        if key in result:
+            paths = result[key] if isinstance(result[key], (list, tuple)) else [result[key]]
+            result[key] = [str(_path(p)) for p in paths]
+    return result
+
+
+def collect(robot, *, teleop=None, output_dir, dataset_format="hdf5", task="",
+            frequency=25, visualize=False, **runtime_options):
+    """Collect demonstrations with robot/teleoperator YAMLs in a prepared environment.
+
+    Call from a script guarded by if __name__ == '__main__' on platforms using
+    multiprocessing spawn. This starts devices and runs until interrupted.
+    """
+    if dataset_format not in {'hdf5', 'lerobotv21', 'lerobotv30'}:
+        raise ValueError('dataset_format must be hdf5, lerobotv21, or lerobotv30')
+    if frequency <= 0:
+        raise ValueError('frequency must be positive')
+    args = ['-r', str(_config(robot, 'robot')), '-o', str(_path(output_dir)),
+            '--dataset-format', dataset_format, '--task', task, '-f', str(frequency)]
+    if teleop is not None:
+        args.extend(['-t', str(_config(teleop, 'teleop'))])
+    if visualize:
+        args.append('--visualize')
+    options = _normalize_options(runtime_options)
+    options.setdefault('runtime', 'current')
+    _run('collect', args, options)
+
+
+def infer(robot, policy, *, output_dir=None, device="cuda", action_manager="basic",
+          publish_rate=25, sensing_rate=20, visualize=False, **runtime_options):
+    """Deploy a Policy checkpoint or RemotePolicy on a configured robot.
+
+    Call from a script guarded by if __name__ == '__main__'. The robot and model
+    dependencies must already be installed in the selected environment.
+    """
+    if isinstance(policy, Policy) and policy.checkpoint is not None:
+        model = str(policy.checkpoint)
+        options = dict(policy.runtime_options)
+    elif isinstance(policy, RemotePolicy):
+        model, options = policy.address, {}
+    elif isinstance(policy, str):
+        model, options = policy, {}
+    else:
+        raise ValueError('infer() requires a saved Policy, RemotePolicy, checkpoint path, or server address')
+    if publish_rate <= 0 or sensing_rate <= 0:
+        raise ValueError('publish_rate and sensing_rate must be positive')
+    args = ['-r', str(_config(robot, 'robot')), '-m', model, '--device', device,
+            '-am', str(_config(action_manager, 'action_manager')),
+            '-pr', str(publish_rate), '-sr', str(sensing_rate)]
+    if output_dir is not None:
+        args.extend(['-o', str(_path(output_dir))])
+    if visualize:
+        args.append('--visualize')
+    options.update(_normalize_options(runtime_options))
+    options.setdefault('runtime', 'current')
+    _run('infer', args, options)
+
+
+def train(policy, dataset, training_config_path="default", *, output_dir,
+          overrides: Mapping | None = None, **runtime_options):
+    """Train synchronously, update policy.checkpoint on success and return artifacts.
+
+    Uses the original policy-specific processors, collators, cache and trainer.
+    Workers use the calling Python environment. Install dependencies with pip/uv
+    before running your script; this API never prepares environments.
+    """
+    if not isinstance(policy, Policy) or not isinstance(dataset, Dataset):
+        raise TypeError("Use load_policy() and load_dataset() handles with train()")
+    output = _path(output_dir)
+    reserved = {"policy", "task", "training_config", "output_dir", "p", "t", "c", "o"}
+    if reserved.intersection(overrides or {}):
+        raise ValueError("Use function parameters for policy, task, training_config and output_dir")
+    args = ["-p", str(policy.config_path), "-t", str(dataset.config_path),
+            "-c", str(_config(training_config_path, "training")), "-o", str(output),
+            *_overrides(overrides)]
+    options = {**policy.runtime_options, **_normalize_options(runtime_options)}
+    if dataset.cache_dir is not None:
+        options["data_cache_dir"] = str(dataset.cache_dir)
+    _run("train", args, options)
+    if not output.is_dir():
+        raise TaskError(f"Training exited successfully but did not create {output}")
+    policy.checkpoint = output
+    # Carry cache settings into evaluation, but not a training-only entrypoint.
+    policy.runtime_options.update({k: v for k, v in options.items() if k != "runtime_manifest"})
+    return TrainingResult(output, policy)

vlastudio/benchmark/aloha/__init__.py

@@ -0,0 +1,143 @@
+"""
+Observation Format
+===============================================================================================================
+sim_transfer_cube_scripted / sim_insertion_scripted
+---------------------------------------------------------------------------------------------------------------
+obs [dict] = {
+    'qpos': ndarray:(14,),
+    'qvel': ndarray:(14,),
+    'env_state': ndarray:(7,),
+    'images': {
+        'top': ndarray:(480, 460, 3),
+        'angle': ndarray:(480, 460, 3),
+        'vis': ndarray:(480, 460, 3),
+    }
+}
+---------------------------------------------------------------------------------------------------------------
+act [ndarray(14,)]
+---------------------------------------------------------------------------------------------------------------
+===============================================================================================================
+Environment for simulated robot bi-manual manipulation, with joint position control
+Action space:      [left_arm_qpos (6),             # absolute joint position
+                    left_gripper_positions (1),    # normalized gripper position (0: close, 1: open)
+                    right_arm_qpos (6),            # absolute joint position
+                    right_gripper_positions (1),]  # normalized gripper position (0: close, 1: open)
+
+Observation space: {"qpos": Concat[ left_arm_qpos (6),         # absolute joint position
+                                    left_gripper_position (1),  # normalized gripper position (0: close, 1: open)
+                                    right_arm_qpos (6),         # absolute joint position
+                                    right_gripper_qpos (1)]     # normalized gripper position (0: close, 1: open)
+                    "qvel": Concat[ left_arm_qvel (6),         # absolute joint velocity (rad)
+                                    left_gripper_velocity (1),  # normalized gripper velocity (pos: opening, neg: closing)
+                                    right_arm_qvel (6),         # absolute joint velocity (rad)
+                                    right_gripper_qvel (1)]     # normalized gripper velocity (pos: opening, neg: closing)
+                    "images": {"main": (480x640x3)}        # h, w, c, dtype='uint8'
+"""
+import os
+import sys
+
+# Configure offscreen rendering before dm_control or MuJoCo is imported.
+if (sys.platform.startswith("linux")
+        and not os.environ.get("DISPLAY")
+        and not os.environ.get("WAYLAND_DISPLAY")):
+    os.environ.setdefault("MUJOCO_GL", "egl")
+
+# Removed problematic numpy import that causes compatibility issues
+from vlastudio.benchmark.base import MetaEnv, MetaAction, MetaObs, MetaPolicy
+from .constants import SIM_TASK_CONFIGS
+from .ee_sim_env import make_ee_sim_env
+from .sim_env import make_sim_env, BOX_POSE
+from .utils import sample_box_pose, sample_insertion_pose
+import numpy as np
+from dataclasses import asdict
+import numpy as np
+import time
+import cv2
+from multiprocessing import current_process
+from loguru import logger
+
+TASK_PROMPT = {
+    'sim_transfer_cube_scripted': 'Transfer the red cube from the right arm to the left arm.',
+    'sim_insertion_scripted': 'Insert the red peg into the blue socket.'
+}
+
+def create_env(config):
+    return AlohaSimEnv(config)
+
+class AlohaSimEnv(MetaEnv):
+    def __init__(self, config, *args):
+        # 初始化env，仅从 config 读取参数
+        self.config = config
+        self.ctrl_space = getattr(self.config, 'ctrl_space', 'joint')
+        self.ctrl_type = getattr(self.config, 'ctrl_type', 'abs')
+        self.camera_ids = getattr(self.config, 'camera_ids', [0])
+        # 统一使用 image_size；同时用于 primary 与 wrist
+        image_size = getattr(self.config, 'image_size', [640, 480])
+        if isinstance(image_size, int):
+            width, height = image_size, image_size
+        elif isinstance(image_size, (list, tuple)):
+            # assume [w, h] or (w, h)
+            width, height = image_size
+        else:
+            raise ValueError("image_size should be list [width, height] or int")
+        self.image_size_primary = (width, height)
+        self.image_size_wrist = (width, height)
+        env = self.create_env()
+        super().__init__(env)
+        
+    def create_env(self):
+        self.task_name = 'transfer_cube' if 'transfer' in self.config.task else 'insertion'
+        self.raw_lang = TASK_PROMPT['sim_transfer_cube_scripted'] if 'transfer' in self.task_name else TASK_PROMPT['sim_insertion_scripted']
+        # step over the environment
+        env = make_sim_env('sim_'+self.task_name)
+        # Get action bounds from action_spec
+        action_spec = env.action_spec()
+        self.min_action = action_spec.minimum
+        self.max_action = action_spec.maximum
+        return env
+        
+    def meta2act(self, maction: MetaAction):
+        # assert maction['ctrl_space']==self.ctrl_space, f"The ctrl_space of MetaAction {maction['ctrl_space']} doesn't match the action space of environment {self.ctrl_space}"
+        # assert maction['ctrl_type']==self.ctrl_type, "Action must be relative action for LIBERO"
+        actions = maction['action'] # (action_dim, )
+        return actions
+        
+    def obs2meta(self, obs):
+        state_joint = np.concatenate([obs["qpos"], ], axis=0).astype(np.float32)
+        # image
+        img_primary = cv2.resize(obs['images']['top'], self.image_size_primary)
+        img_wrist_left = cv2.resize(obs['images']['left_wrist'], self.image_size_wrist)
+        img_wrist_right = cv2.resize(obs['images']['right_wrist'], self.image_size_wrist)
+        all_imgs = [img_primary]
+        if '1' in self.camera_ids or 1 in self.camera_ids:
+            all_imgs = all_imgs + [img_wrist_left, ]
+        if '2' in self.camera_ids or 2 in self.camera_ids:
+            all_imgs = all_imgs + [img_wrist_right, ]
+        image = np.stack(all_imgs)
+        image = image.transpose(0, 3, 1, 2)
+        return MetaObs(state=state_joint, image=image, raw_lang=self.raw_lang)
+    
+    def step(self, *args, **kwargs):
+        act = self.meta2act(*args, **kwargs)
+        ts = self.env.step(act)
+        results = [ts.observation, ts.reward, self.env.task.max_reward==ts.reward, ts]
+        self.prev_obs = results[0] = self.obs2meta(results[0])
+        if isinstance(results[0], MetaObs): results[0] = asdict(results[0])
+        return tuple(results)
+
+    def reset(self):
+        pid = current_process().pid  # 获取当前进程 ID
+        seed = (pid * 1000 + time.time_ns()) % (2**32)  # 基于时间戳生成种子
+        np.random.seed(seed)
+        # global BOX_POSE
+        if 'transfer' in self.task_name:
+            BOX_POSE[0] = sample_box_pose()  # used in sim reset
+        elif 'insertion' in self.task_name:
+            BOX_POSE[0] = np.concatenate(sample_insertion_pose())  # used in sim reset
+        ts = self.env.reset()
+        self.prev_obs = self.obs2meta(ts.observation)
+        return self.prev_obs
+    
+    def close(self):
+        self.env.close()
+    

vlastudio/benchmark/aloha/assets/bimanual_viperx_ee_insertion.xml

@@ -0,0 +1,59 @@
+<mujoco>
+    <include file="scene.xml"/>
+    <include file="vx300s_dependencies.xml"/>
+
+    <equality>
+        <weld body1="mocap_left" body2="vx300s_left/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
+        <weld body1="mocap_right" body2="vx300s_right/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
+    </equality>
+
+
+    <worldbody>
+        <include file="vx300s_left.xml" />
+        <include file="vx300s_right.xml" />
+
+        <body mocap="true" name="mocap_left" pos="0.095 0.50 0.425">
+            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_left_site1" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_left_site2" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_left_site3" rgba="1 0 0 1"/>
+        </body>
+        <body mocap="true" name="mocap_right" pos="-0.095 0.50 0.425">
+            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_right_site1" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_right_site2" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_right_site3" rgba="1 0 0 1"/>
+        </body>
+
+        <body name="peg" pos="0.2 0.5 0.05">
+            <joint name="red_peg_joint" type="free" frictionloss="0.01" />
+            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" />
+        </body>
+
+        <body name="socket" pos="-0.2 0.5 0.05">
+            <joint name="blue_socket_joint" type="free" frictionloss="0.01" />
+            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+<!--            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg_ref" rgba="1 0 0 1" />-->
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" />
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" />
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" />
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" />
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" />
+        </body>
+
+    </worldbody>
+
+    <actuator>
+        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200"  user="1"/>
+        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200"  user="1"/>
+
+        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200"  user="1"/>
+        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200"  user="1"/>
+
+    </actuator>
+
+    <keyframe>
+        <key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0.2 0.5 0.05 1 0 0 0  -0.2 0.5 0.05 1 0 0 0"/>
+    </keyframe>
+
+
+</mujoco>

vlastudio/benchmark/aloha/assets/bimanual_viperx_ee_transfer_cube.xml

@@ -0,0 +1,48 @@
+<mujoco>
+    <include file="scene.xml"/>
+    <include file="vx300s_dependencies.xml"/>
+
+    <equality>
+        <weld body1="mocap_left" body2="vx300s_left/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
+        <weld body1="mocap_right" body2="vx300s_right/gripper_link" solref="0.01 1" solimp=".25 .25 0.001" />
+    </equality>
+
+
+    <worldbody>
+        <include file="vx300s_left.xml" />
+        <include file="vx300s_right.xml" />
+
+        <body mocap="true" name="mocap_left" pos="0.095 0.50 0.425">
+            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_left_site1" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_left_site2" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_left_site3" rgba="1 0 0 1"/>
+        </body>
+        <body mocap="true" name="mocap_right" pos="-0.095 0.50 0.425">
+            <site pos="0 0 0" size="0.003 0.003 0.03" type="box" name="mocap_right_site1" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.003 0.03 0.003" type="box" name="mocap_right_site2" rgba="1 0 0 1"/>
+            <site pos="0 0 0" size="0.03 0.003 0.003" type="box" name="mocap_right_site3" rgba="1 0 0 1"/>
+        </body>
+
+        <body name="box" pos="0.2 0.5 0.05">
+            <joint name="red_box_joint" type="free" frictionloss="0.01" />
+            <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+            <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.02 0.02 0.02" type="box" name="red_box" rgba="1 0 0 1" />
+        </body>
+
+    </worldbody>
+
+    <actuator>
+        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_left/left_finger" kp="200"  user="1"/>
+        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_left/right_finger" kp="200"  user="1"/>
+
+        <position ctrllimited="true" ctrlrange="0.021 0.057" joint="vx300s_right/left_finger" kp="200"  user="1"/>
+        <position ctrllimited="true" ctrlrange="-0.057 -0.021" joint="vx300s_right/right_finger" kp="200"  user="1"/>
+
+    </actuator>
+
+    <keyframe>
+        <key qpos="0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0 -0.96 1.16 0 -0.3 0 0.024 -0.024  0.2 0.5 0.05 1 0 0 0"/>
+    </keyframe>
+
+
+</mujoco>