commandagi 0.3.0__py3-none-any.whl

commandagi/lerobot.py

@@ -0,0 +1,211 @@
+"""Record CommandAGI sim episodes as a LeRobot-style dataset — observation (camera + joint state),
+action, and timestamps per step, written to disk in the LeRobot on-disk layout (parquet per episode
++ frames + meta/info.json) so it loads into imitation/VLA training pipelines.
+
+Two ways to use it:
+
+1. Wrap a Gym env — every ``step`` is recorded automatically::
+
+       from commandagi.gym_env import CommandAGIEnv
+       from commandagi.lerobot import RecordingWrapper
+       env = RecordingWrapper(CommandAGIEnv(world), root="./datasets/pick", task="pick up the cube")
+       obs, _ = env.reset()
+       for _ in range(200):
+           obs, *_ = env.step(policy(obs))
+       env.save()            # flush the dataset to ./datasets/pick
+
+2. Hand-drive and log directly with :class:`EpisodeRecorder` (e.g. while teleoperating in Sim Studio).
+
+This writes the LeRobot v2.1/3.0 directory shape (``meta/info.json``, ``data/chunk-000/episode_*.parquet``,
+``images/observation.image/episode_*/frame_*.jpg``). Requires ``pyarrow`` + ``numpy`` + ``pillow``
+(``pip install "commandagi[lerobot]"``). Stats consolidation (``meta/stats.json``) is left to the
+``lerobot`` tooling; everything needed to compute it is on disk.
+"""
+from __future__ import annotations
+
+import io
+import json
+import time
+from pathlib import Path
+from typing import Any, Optional
+
+try:
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    from PIL import Image
+except ImportError as e:  # pragma: no cover
+    raise ImportError("commandagi.lerobot needs `pyarrow`, `numpy`, `pillow` — pip install 'commandagi[lerobot]'") from e
+
+CODEBASE_VERSION = "v2.1"
+
+
+class EpisodeRecorder:
+    """Accumulates steps for one or more episodes and writes a LeRobot-layout dataset.
+
+    Each step is ``(state, action, image)`` where ``state``/``action`` are 1-D float arrays and
+    ``image`` is an HxWx3 uint8 array (optional). Call :meth:`add` per step, :meth:`end_episode` to
+    close the current episode, and :meth:`save` to write everything out.
+    """
+
+    def __init__(self, root: str, *, fps: float = 10.0, task: str = "", image_key: str = "observation.image"):
+        self.root = Path(root)
+        self.fps = float(fps)
+        self.task = task
+        self.image_key = image_key
+        self._ep: list[dict[str, Any]] = []  # steps in the current episode
+        self._episodes: list[list[dict[str, Any]]] = []  # completed episodes
+        self._state_dim: Optional[int] = None
+        self._action_dim: Optional[int] = None
+        self._img_shape: Optional[tuple[int, int, int]] = None
+        self._t0 = time.time()
+
+    def start_episode(self) -> None:
+        if self._ep:
+            self.end_episode()
+        self._ep = []
+        self._t0 = time.time()
+
+    def add(self, state, action, image=None, *, reward: float = 0.0) -> None:
+        state = np.asarray(state, dtype=np.float32).reshape(-1)
+        action = np.asarray(action, dtype=np.float32).reshape(-1)
+        self._state_dim = self._state_dim or state.shape[0]
+        self._action_dim = self._action_dim or action.shape[0]
+        step: dict[str, Any] = {
+            "observation.state": state,
+            "action": action,
+            "timestamp": time.time() - self._t0,
+            "reward": float(reward),
+        }
+        if image is not None:
+            img = np.asarray(image, dtype=np.uint8)
+            self._img_shape = self._img_shape or tuple(img.shape)  # type: ignore[assignment]
+            step["_image"] = img
+        self._ep.append(step)
+
+    def end_episode(self) -> None:
+        if self._ep:
+            self._episodes.append(self._ep)
+            self._ep = []
+
+    def save(self) -> Path:
+        """Write the dataset (all completed episodes + the current one) to disk; returns the root."""
+        self.end_episode()
+        if not self._episodes:
+            raise ValueError("nothing to save — record at least one step")
+        (self.root / "data" / "chunk-000").mkdir(parents=True, exist_ok=True)
+        (self.root / "meta").mkdir(parents=True, exist_ok=True)
+
+        total_frames = 0
+        global_index = 0
+        for ep_idx, steps in enumerate(self._episodes):
+            cols: dict[str, list] = {
+                "observation.state": [], "action": [], "timestamp": [], "reward": [],
+                "frame_index": [], "episode_index": [], "index": [], "next.done": [],
+            }
+            has_image = any("_image" in s for s in steps)
+            if has_image:
+                cols[self.image_key] = []
+                img_dir = self.root / "images" / self.image_key / f"episode_{ep_idx:06d}"
+                img_dir.mkdir(parents=True, exist_ok=True)
+            for f_idx, s in enumerate(steps):
+                cols["observation.state"].append(s["observation.state"].tolist())
+                cols["action"].append(s["action"].tolist())
+                cols["timestamp"].append(float(s["timestamp"]))
+                cols["reward"].append(float(s["reward"]))
+                cols["frame_index"].append(f_idx)
+                cols["episode_index"].append(ep_idx)
+                cols["index"].append(global_index)
+                cols["next.done"].append(f_idx == len(steps) - 1)
+                if has_image:
+                    rel = f"images/{self.image_key}/episode_{ep_idx:06d}/frame_{f_idx:06d}.jpg"
+                    if "_image" in s:
+                        Image.fromarray(s["_image"]).save(self.root / rel, format="JPEG", quality=90)
+                    cols[self.image_key].append(rel)
+                global_index += 1
+            total_frames += len(steps)
+            pq.write_table(pa.table(cols), self.root / "data" / "chunk-000" / f"episode_{ep_idx:06d}.parquet")
+
+        self._write_info(total_frames)
+        self._write_episodes_meta()
+        self._write_tasks_meta()
+        return self.root
+
+    def _features(self) -> dict:
+        feats: dict[str, Any] = {
+            "observation.state": {"dtype": "float32", "shape": [self._state_dim or 0], "names": None},
+            "action": {"dtype": "float32", "shape": [self._action_dim or 0], "names": None},
+            "timestamp": {"dtype": "float32", "shape": [1], "names": None},
+            "reward": {"dtype": "float32", "shape": [1], "names": None},
+            "frame_index": {"dtype": "int64", "shape": [1], "names": None},
+            "episode_index": {"dtype": "int64", "shape": [1], "names": None},
+            "index": {"dtype": "int64", "shape": [1], "names": None},
+            "next.done": {"dtype": "bool", "shape": [1], "names": None},
+        }
+        if self._img_shape is not None:
+            feats[self.image_key] = {"dtype": "image", "shape": list(self._img_shape), "names": ["height", "width", "channel"]}
+        return feats
+
+    def _write_info(self, total_frames: int) -> None:
+        info = {
+            "codebase_version": CODEBASE_VERSION,
+            "robot_type": "commandagi-sim",
+            "total_episodes": len(self._episodes),
+            "total_frames": total_frames,
+            "total_tasks": 1,
+            "total_chunks": 1,
+            "chunks_size": 1000,
+            "fps": self.fps,
+            "splits": {"train": f"0:{len(self._episodes)}"},
+            "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet",
+            "features": self._features(),
+        }
+        (self.root / "meta" / "info.json").write_text(json.dumps(info, indent=2))
+
+    def _write_episodes_meta(self) -> None:
+        lines = []
+        for ep_idx, steps in enumerate(self._episodes):
+            lines.append(json.dumps({"episode_index": ep_idx, "tasks": [self.task] if self.task else [], "length": len(steps)}))
+        (self.root / "meta" / "episodes.jsonl").write_text("\n".join(lines) + "\n")
+
+    def _write_tasks_meta(self) -> None:
+        (self.root / "meta" / "tasks.jsonl").write_text(json.dumps({"task_index": 0, "task": self.task}) + "\n")
+
+
+class RecordingWrapper:
+    """Wrap a :class:`commandagi.gym_env.CommandAGIEnv` so every ``reset``/``step`` is logged.
+
+    The observation's ``qpos``+``qvel`` become ``observation.state``; ``image`` (if present) is the
+    camera frame; the action passed to ``step`` is recorded verbatim. Call :meth:`save` when done.
+    """
+
+    def __init__(self, env, root: str, *, task: str = "", fps: Optional[float] = None):
+        self.env = env
+        self.recorder = EpisodeRecorder(root, fps=fps or (1.0 / getattr(env, "control_dt", 0.1)), task=task)
+        self.action_space = env.action_space
+        self.observation_space = env.observation_space
+
+    @staticmethod
+    def _state(obs: dict) -> np.ndarray:
+        return np.concatenate([np.asarray(obs.get("qpos", []), dtype=np.float32), np.asarray(obs.get("qvel", []), dtype=np.float32)])
+
+    def reset(self, **kw):
+        obs, info = self.env.reset(**kw)
+        self.recorder.start_episode()
+        self._last_obs = obs
+        return obs, info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        # Record the transition: the observation we acted on + the action + the reward received.
+        self.recorder.add(self._state(self._last_obs), action, self._last_obs.get("image"), reward=reward)
+        self._last_obs = obs
+        if terminated or truncated:
+            self.recorder.end_episode()
+        return obs, reward, terminated, truncated, info
+
+    def save(self):
+        return self.recorder.save()
+
+    def close(self):
+        self.env.close()

commandagi-0.3.0.dist-info/METADATA

@@ -0,0 +1,159 @@
+Metadata-Version: 2.5
+Name: commandagi
+Version: 0.3.0
+Summary: Python SDK for CommandAGI — launch cloud computers and 3D robot simulations and control them programmatically.
+Project-URL: Homepage, https://commandagi.com
+Project-URL: Documentation, https://github.com/CommandAGI/commandagi/blob/main/docs/ROBOT_DEVELOPER_API.md
+Author: CommandAGI
+License: MIT
+Keywords: agents,computer-use,physics,reinforcement-learning,robotics,simulation
+Requires-Python: >=3.9
+Requires-Dist: requests>=2.28
+Requires-Dist: websocket-client>=1.5
+Provides-Extra: gym
+Requires-Dist: gymnasium>=0.29; extra == 'gym'
+Requires-Dist: numpy>=1.23; extra == 'gym'
+Requires-Dist: pillow>=9.0; extra == 'gym'
+Provides-Extra: lerobot
+Requires-Dist: numpy>=1.23; extra == 'lerobot'
+Requires-Dist: pillow>=9.0; extra == 'lerobot'
+Requires-Dist: pyarrow>=14.0; extra == 'lerobot'
+Provides-Extra: vision
+Requires-Dist: numpy>=1.23; extra == 'vision'
+Requires-Dist: pillow>=9.0; extra == 'vision'
+Description-Content-Type: text/markdown
+
+# CommandAGI Python SDK
+
+Launch real cloud **computers** and **3D robot simulations** and control them from Python — stream
+the robot's camera, send actions, run episodes. No agent required: you drive.
+
+```bash
+pip install commandagi          # + `pip install commandagi[vision]` for numpy frames
+```
+
+## Robot testing in a 3D world
+
+```python
+from commandagi import CommandAGI
+
+cagi = CommandAGI(api_key="cagi_...")          # or set COMMANDAGI_API_KEY
+
+with cagi.launch("simulation/warehouse") as world:
+    obs = world.observe()                      # JPEG bytes from the robot's head camera
+    for _ in range(20):
+        obs = world.step("turn", dir="left")   # act, then get the next frame
+    world.reset()                              # robot back to the episode start
+# leaving the block stops the world and releases the cloud VM
+```
+
+`launch()` provisions a real GCE VM running a 3D physics world, waits until it's streaming, and gives
+you a `World`. Built-in scenes: `simulation/warehouse`, `simulation/house-on-fire`,
+`simulation/school` (a mobile robot in each).
+
+### The control vocabulary
+
+| World kind  | actions                                                                   |
+| ----------- | ------------------------------------------------------------------------- |
+| robot / sim | `move(speed)`, `back(speed)`, `turn(dir, rate)`, `stop`, `reset`          |
+| computer    | `click(x, y)`, `type(text)`, `key(key)`, `move(x, y)`, `scroll(x, y, dy)` |
+
+```python
+world.act("move", speed=0.8)        # fire-and-forget
+obs = world.step("move", speed=0.8) # act + return the next observation (settles 0.8s)
+obs = world.observe(fresh=True)     # wait for a frame newer than now
+arr = world.observe_array()         # HxWx3 uint8 numpy (needs commandagi[vision])
+for frame in world.stream():        # live generator of frames
+    ...
+```
+
+## Simulator instances (morphology-agnostic robots)
+
+The simulator is **morphology-agnostic**: a robot is just a set of named actuators and sites, driven
+by one small **generic** control vocabulary — no `drive`/`gripper`, just `ctrl` / `actuator` / `ik`
+/ `trajectory` / `describe`. Spin up your own instance, choose who can watch or add robots, and
+populate it with one or many robots on a single session.
+
+```python
+from commandagi import CommandAGI
+
+cagi = CommandAGI(api_key="cagi_...")
+
+sim = cagi.launch_sim(scene="the-matrix", visibility="private", title="demo")
+print("instance:", sim.id, "session:", sim.session_id)
+
+# Who can do what:
+sim.grant("user_teammate", capability="viewer")    # may watch the stream
+sim.grant("user_buddy",    capability="operator")  # may also launch robots into the world
+
+# Add robots (each becomes a embodiment on sim.session_id):
+rover = sim.join_robot(kind="rover")   # -> {robotId, embodimentId, sessionId}
+arm   = sim.join_robot(kind="arm")
+
+cagi.sims()             # list instances visible to you
+cagi.get_sim(sim.id)    # rehydrate a SimInstance
+sim.view()              # instance metadata + attached embodiments
+sim.stop()              # release it (or use `with cagi.launch_sim(...) as sim:`)
+```
+
+### Generic robot control
+
+`World` exposes the morphology-agnostic vocabulary (address a specific robot in a multi-robot embodiment
+with `robot_id`):
+
+```python
+world = cagi.connect_world(sim.session_id, rover["embodimentId"], kind="robot")
+
+desc = world.describe()                                   # actuators, sites, objects (best-effort)
+world.ctrl({"left_wheel": 1.0, "right_wheel": 1.0})       # set actuator targets directly
+world.actuator("left_wheel", 0.0)                          # one named actuator
+world.ik(target=[0.3, 0.0, 0.4], site="ee", relative=False)  # inverse kinematics to a point
+world.trajectory([{"left_wheel": 1.0}, {"left_wheel": 0.0}])  # follow waypoints
+frame = world.observe()                                    # camera frame, as before
+```
+
+> `describe()` is best-effort: the runtime answers a describe request over the session channel, but
+> there is currently no synchronous describe HTTP endpoint — if nothing echoes back it returns `{}`,
+> and the autonomous agent also obtains descriptions server-side via `/agent/robot-act`.
+
+## Autonomous agents over many robots
+
+One agent can drive **many** robots in a single session. `RobotAgent` loops perceive → reason → act:
+each step it gathers every embodiment's description + a fresh camera frame, calls `/agent/robot-act` with
+all embodiments, and applies the returned tool calls (`ctrl`/`actuator`/`ik`/`trajectory`) back to the
+addressed embodiment — until a `done` call or `max_steps`.
+
+```python
+from commandagi import CommandAGI
+from commandagi.agent import RobotAgent, attach_robots
+
+cagi = CommandAGI(api_key="cagi_...")
+sim = cagi.launch_sim(scene="warehouse")
+
+embodiments = attach_robots(cagi, sim, kinds=["rover", "arm"])   # two robots, one session
+
+with RobotAgent(cagi, sim.session_id, embodiments, goal="bring the red box to the arm") as agent:
+    result = agent.run(max_steps=25)        # blocks; prints reasoning + applied calls each step
+print("done:", result["done"], "in", result["steps"], "steps")
+
+sim.stop()
+```
+
+A full runnable script lives in [`examples/sim_agent.py`](examples/sim_agent.py).
+
+## Computers too
+
+```python
+with cagi.launch("computer/software-engineer") as pc:
+    pc.act("type", text="hello")
+    pc.act("key", key="Return")
+    screenshot = pc.observe()       # PNG bytes of the live Ubuntu desktop
+```
+
+## Auth
+
+Create an API key with an `operator` scope (dashboard → API keys, or `POST /me/api-keys`). Pass it to
+`CommandAGI(api_key=...)` or set `COMMANDAGI_API_KEY`. Point at another environment with
+`COMMANDAGI_BASE_URL` (e.g. `https://api-dev.commandagi.com`).
+
+Full HTTP + WebSocket reference (what the SDK wraps): [`docs/platform/ROBOT_DEVELOPER_API.md`](../../docs/platform/ROBOT_DEVELOPER_API.md).

commandagi-0.3.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+commandagi/__init__.py,sha256=aF679wR7Fa0-7Xg8swPgQZ0PKB7LlyMgiwEiWUWknw0,697
+commandagi/agent.py,sha256=hmnyWYQ9Ma7gZYQVNmXQKv-3YC70FQs_3KjpWXkSaq8,7737
+commandagi/bridge.py,sha256=yS225aHQtYTsLEAz8nAxwG-qqUpidtqcbbh-deflwfU,5260
+commandagi/client.py,sha256=LeWwaqe7zjX86Onr_OtHCgxgNiIryxEPIu1RAugJHpQ,24608
+commandagi/gym_env.py,sha256=a9jKDGr5qI-e49E5Mm-EvGMRlTRvPQKErYsg2yCWv5s,6405
+commandagi/lerobot.py,sha256=b-hP3MXnGLBuOoE0DsWvS_iLYAx3-V3oojttHRSKJdc,9488
+commandagi-0.3.0.dist-info/METADATA,sha256=rAVYv4ChUdWxbutZ2DP9dokFJGZCe0LtmBF6qGySobY,6860
+commandagi-0.3.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+commandagi-0.3.0.dist-info/RECORD,,

commandagi-0.3.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any