diff --git a/.gitignore b/.gitignore index 8f4c8e1..b46f014 100644 --- a/.gitignore +++ b/.gitignore @@ -175,7 +175,7 @@ outputs/ .cache/* datasets -20250901 +2025* s100 huggingface_models diff --git a/Makefile b/Makefile index 5bfbe76..dfa5807 100644 --- a/Makefile +++ b/Makefile @@ -1,180 +1,109 @@ -# 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. +VERSION := 20251229-release +REMOTE_VERSION := v1.1.0 +PROXY := http://192.168.16.68:18000 +REMOTE_REPO := ccr-29eug8s3-pub.cnc.bj.baidubce.com/public +CUDA_DEVICES := 0 -.PHONY: tests +# Default target +.PHONY: all +all: help -PYTHON_PATH := $(shell which python) +.PHONY: build +build: build-merge build-train build-infer build-visualize + +.PHONY: publish +push: build push-merge push-train push-infer push-visualize clean -# If uv is installed and a virtual environment exists, use it -UV_CHECK := $(shell command -v uv) -ifneq ($(UV_CHECK),) - PYTHON_PATH := $(shell .venv/bin/python) -endif +.PHONY: merge train infer visualize +merge: build-merge run-merge +train: build-train run-train +infer: build-infer run-infer +visualize: build-visualize run-visualize -export PATH := $(dir $(PYTHON_PATH)):$(PATH) +# --- Merge Image --- +.PHONY: build-merge run-merge push-merge +build-merge: + docker build -t dcloud/lerobot-merge:$(VERSION) -f docker/Dockerfile.merge . --build-arg https_proxy=$(PROXY) -DEVICE ?= cpu +run-merge: + docker run -it --rm \ + -v $(CURDIR)/datasets:/workspace/embolab/inputs \ + -v $(CURDIR)/docker/inputs:/workspace/embolab/outputs \ + -v $(CURDIR)/docker/merge_task.json:/workspace/embolab/params/build_task.json \ + -e HF_HUB_OFFLINE=1 -e HF_DATASET_OFFLINE=1 \ + --shm-size=120G \ + dcloud/lerobot-merge:$(VERSION) -build-user: - docker build -f docker/Dockerfile.user -t lerobot-user . +push-merge: + docker tag dcloud/lerobot-merge:$(VERSION) $(REMOTE_REPO)/lerobot-merge:$(REMOTE_VERSION) + docker push $(REMOTE_REPO)/lerobot-merge:$(REMOTE_VERSION) -build-internal: - docker build -f docker/Dockerfile.internal -t lerobot-internal . +# --- Training Image --- +.PHONY: build-train run-train push-train +build-train: + docker build -t dcloud/lerobot-train:$(VERSION) -f docker/Dockerfile.train . --build-arg https_proxy=$(PROXY) -test-end-to-end: - ${MAKE} DEVICE=$(DEVICE) test-act-ete-train - ${MAKE} DEVICE=$(DEVICE) test-act-ete-train-resume - ${MAKE} DEVICE=$(DEVICE) test-act-ete-eval - ${MAKE} DEVICE=$(DEVICE) test-diffusion-ete-train - ${MAKE} DEVICE=$(DEVICE) test-diffusion-ete-eval - ${MAKE} DEVICE=$(DEVICE) test-tdmpc-ete-train - ${MAKE} DEVICE=$(DEVICE) test-tdmpc-ete-eval - ${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-train - ${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval +run-train: + docker run -it --rm --gpus '"device=$(CUDA_DEVICES)"' \ + -v $(CURDIR)/docker/inputs:/workspace/embolab/inputs \ + -v $(CURDIR)/docker/outputs:/workspace/embolab/outputs/ \ + -v $(CURDIR)/docker/train_task.json:/workspace/embolab/params/build_task.json \ + -e HF_HUB_OFFLINE=1 -e HF_DATASET_OFFLINE=1 \ + --shm-size=120G \ + dcloud/lerobot-train:$(VERSION) -test-act-ete-train: - python -m lerobot.scripts.train \ - --policy.type=act \ - --policy.dim_model=64 \ - --policy.n_action_steps=20 \ - --policy.chunk_size=20 \ - --policy.device=$(DEVICE) \ - --policy.push_to_hub=false \ - --env.type=aloha \ - --env.episode_length=5 \ - --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \ - --dataset.image_transforms.enable=true \ - --dataset.episodes="[0]" \ - --batch_size=2 \ - --steps=4 \ - --eval_freq=2 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 \ - --save_freq=2 \ - --save_checkpoint=true \ - --log_freq=1 \ - --wandb.enable=false \ - --output_dir=tests/outputs/act/ +push-train: + docker tag dcloud/lerobot-train:$(VERSION) $(REMOTE_REPO)/lerobot-train:$(REMOTE_VERSION) + docker push $(REMOTE_REPO)/lerobot-train:$(REMOTE_VERSION) -test-act-ete-train-resume: - python -m lerobot.scripts.train \ - --config_path=tests/outputs/act/checkpoints/000002/pretrained_model/train_config.json \ - --resume=true +# --- Online Inference Image --- +.PHONY: build-infer run-infer push-infer +build-infer: + docker build -t dcloud/lerobot-infer:$(VERSION) -f docker/Dockerfile.infer . --build-arg https_proxy=$(PROXY) -test-act-ete-eval: - python -m lerobot.scripts.eval \ - --policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \ - --policy.device=$(DEVICE) \ - --env.type=aloha \ - --env.episode_length=5 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 +run-infer: + docker run -it --rm --gpus '"device=$(CUDA_DEVICES)"' \ + -v $(CURDIR)/docker/outputs:/workspace/embolab/outputs \ + -v $(CURDIR)/docker/infer_task.json:/workspace/embolab/params/build_task.json \ + -e HF_HUB_OFFLINE=1 -e HF_DATASET_OFFLINE=1 \ + --shm-size=128G \ + dcloud/lerobot-infer:$(VERSION) -test-diffusion-ete-train: - python -m lerobot.scripts.train \ - --policy.type=diffusion \ - --policy.down_dims='[64,128,256]' \ - --policy.diffusion_step_embed_dim=32 \ - --policy.num_inference_steps=10 \ - --policy.device=$(DEVICE) \ - --policy.push_to_hub=false \ - --env.type=pusht \ - --env.episode_length=5 \ - --dataset.repo_id=lerobot/pusht \ - --dataset.image_transforms.enable=true \ - --dataset.episodes="[0]" \ - --batch_size=2 \ - --steps=2 \ - --eval_freq=2 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 \ - --save_checkpoint=true \ - --save_freq=2 \ - --log_freq=1 \ - --wandb.enable=false \ - --output_dir=tests/outputs/diffusion/ +push-infer: + docker tag dcloud/lerobot-infer:$(VERSION) $(REMOTE_REPO)/lerobot-infer:$(REMOTE_VERSION) + docker push $(REMOTE_REPO)/lerobot-infer:$(REMOTE_VERSION) -test-diffusion-ete-eval: - python -m lerobot.scripts.eval \ - --policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \ - --policy.device=$(DEVICE) \ - --env.type=pusht \ - --env.episode_length=5 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 +# --- Visualize Image --- +.PHONY: build-visualize run-visualize push-visualize +build-visualize: + docker build -t dcloud/lerobot-visualize:$(VERSION) -f docker/Dockerfile.visualize . --build-arg https_proxy=$(PROXY) -test-tdmpc-ete-train: - python -m lerobot.scripts.train \ - --policy.type=tdmpc \ - --policy.device=$(DEVICE) \ - --policy.push_to_hub=false \ - --env.type=xarm \ - --env.task=XarmLift-v0 \ - --env.episode_length=5 \ - --dataset.repo_id=lerobot/xarm_lift_medium \ - --dataset.image_transforms.enable=true \ - --dataset.episodes="[0]" \ - --batch_size=2 \ - --steps=2 \ - --eval_freq=2 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 \ - --save_checkpoint=true \ - --save_freq=2 \ - --log_freq=1 \ - --wandb.enable=false \ - --output_dir=tests/outputs/tdmpc/ +run-visualize: + docker run -it --rm \ + -v $(CURDIR)/datasets:/workspace/embolab/datasets \ + -e HF_HUB_OFFLINE=1 -e HF_DATASET_OFFLINE=1 \ + -p 29876:9876 \ + --shm-size=120G \ + dcloud/lerobot-visualize:$(VERSION) -test-tdmpc-ete-eval: - python -m lerobot.scripts.eval \ - --policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \ - --policy.device=$(DEVICE) \ - --env.type=xarm \ - --env.episode_length=5 \ - --env.task=XarmLift-v0 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 +push-visualize: + docker tag dcloud/lerobot-visualize:$(VERSION) $(REMOTE_REPO)/lerobot-visualize:$(REMOTE_VERSION) + docker push $(REMOTE_REPO)/lerobot-visualize:$(REMOTE_VERSION) +# --- Cleanup --- +.PHONY: clean +clean: + -docker images -f "dangling=true" -q | xargs -r docker rmi + -docker images | grep lerobot | grep -v $(VERSION) | awk '{print $$1":"$$2}' | xargs -r docker rmi -test-smolvla-ete-train: - python -m lerobot.scripts.train \ - --policy.type=smolvla \ - --policy.n_action_steps=20 \ - --policy.chunk_size=20 \ - --policy.device=$(DEVICE) \ - --policy.push_to_hub=false \ - --env.type=aloha \ - --env.episode_length=5 \ - --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \ - --dataset.image_transforms.enable=true \ - --dataset.episodes="[0]" \ - --batch_size=2 \ - --steps=4 \ - --eval_freq=2 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 \ - --save_freq=2 \ - --save_checkpoint=true \ - --log_freq=1 \ - --wandb.enable=false \ - --output_dir=tests/outputs/smolvla/ - -test-smolvla-ete-eval: - python -m lerobot.scripts.eval \ - --policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \ - --policy.device=$(DEVICE) \ - --env.type=aloha \ - --env.episode_length=5 \ - --eval.n_episodes=1 \ - --eval.batch_size=1 +# --- Help --- +.PHONY: help +help: + @echo "Available commands:" + @echo " make build Build all Docker images" + @echo " make push Push all Docker images to remote repository" + @echo " make merge Build and run the merge image" + @echo " make train Build and run the training image" + @echo " make infer Build and run the inference image" + @echo " make visualize Build and run the visualization image" + @echo " make clean Clean up dangling Docker images" \ No newline at end of file diff --git a/docker/Dockerfile.infer b/docker/Dockerfile.infer new file mode 100644 index 0000000..187de8c --- /dev/null +++ b/docker/Dockerfile.infer @@ -0,0 +1,72 @@ +# 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. + +# This Dockerfile is designed for a lerobot user who wants to +# experiment with the project. It starts from an Python Slim base image. + +# docker build -f docker/Dockerfile.user -t lerobot-user . +# docker run -it --rm lerobot-user + +# Configure the base image +ARG PYTHON_VERSION=3.10 +FROM python:${PYTHON_VERSION}-slim + +# Configure environment variables +ENV DEBIAN_FRONTEND=noninteractive \ + PATH=/lerobot/.venv/bin:$PATH + +# Install system dependencies and uv (as root) +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential git curl ffmpeg \ + && curl -LsSf https://astral.sh/uv/install.sh | sh \ + && mv /root/.local/bin/uv /usr/local/bin/uv \ + && useradd --create-home --shell /bin/bash user_lerobot \ + && usermod -aG sudo user_lerobot \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Create application directory and set permissions +RUN git clone https://github.com/huggingface/lerobot.git -b v0.3.3 /lerobot +WORKDIR /lerobot +RUN chown -R user_lerobot:user_lerobot /lerobot + +ADD ./huggingface_models /home/user_lerobot/.cache/huggingface +RUN chown -R user_lerobot:user_lerobot /home/user_lerobot/.cache + + +# Switch to the non-root user +USER user_lerobot + +# Environment variables for the testing +ENV HOME=/home/user_lerobot \ + HF_HOME=/home/user_lerobot/.cache/huggingface \ + HF_LEROBOT_HOME=/home/user_lerobot/.cache/huggingface/lerobot \ + TORCH_HOME=/home/user_lerobot/.cache/torch \ + TRITON_CACHE_DIR=/home/user_lerobot/.cache/triton + +# Create the virtual environment +# We use a virtual environment inside the container—even though the container itself \ +# provides isolation—to closely resemble local development and allow users to \ +# run other Python projects in the same container without dependency conflicts. +RUN uv venv + +# Install Python dependencies for caching +COPY --chown=user_lerobot:user_lerobot pyproject.toml README.md MANIFEST.in ./ +# COPY --chown=user_lerobot:user_lerobot src/ src/ +RUN uv pip install --no-cache ".[smolvla]" +RUN uv pip install pyzmq msgpack msgpack_numpy zstandard + +# Set the default command - Online Inference mode +ADD docker/infer.py /workspace/infer.py +ADD docker/cloud_helper.py /workspace/cloud_helper.py +CMD ["python", "/workspace/infer.py"] diff --git a/docker/Dockerfile.visualize b/docker/Dockerfile.visualize new file mode 100644 index 0000000..ebe31a1 --- /dev/null +++ b/docker/Dockerfile.visualize @@ -0,0 +1,70 @@ +# 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. + +# This Dockerfile is designed for a lerobot user who wants to +# experiment with the project. It starts from an Python Slim base image. + +# docker build -f docker/Dockerfile.user -t lerobot-user . +# docker run -it --rm lerobot-user + +# Configure the base image +ARG PYTHON_VERSION=3.10 +FROM python:${PYTHON_VERSION}-slim + +# Configure environment variables +ENV DEBIAN_FRONTEND=noninteractive \ + PATH=/lerobot/.venv/bin:$PATH + +# Install system dependencies and uv (as root) +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential git curl ffmpeg \ + && curl -LsSf https://astral.sh/uv/install.sh | sh \ + && mv /root/.local/bin/uv /usr/local/bin/uv \ + && useradd --create-home --shell /bin/bash user_lerobot \ + && usermod -aG sudo user_lerobot \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Create application directory and set permissions +RUN git clone https://github.com/huggingface/lerobot.git -b v0.3.3 /lerobot +WORKDIR /lerobot +RUN chown -R user_lerobot:user_lerobot /lerobot + +ADD ./huggingface_models /home/user_lerobot/.cache/huggingface +RUN chown -R user_lerobot:user_lerobot /home/user_lerobot/.cache + + +# Switch to the non-root user +USER user_lerobot + +# Environment variables for the testing +ENV HOME=/home/user_lerobot \ + HF_HOME=/home/user_lerobot/.cache/huggingface \ + HF_LEROBOT_HOME=/home/user_lerobot/.cache/huggingface/lerobot \ + TORCH_HOME=/home/user_lerobot/.cache/torch \ + TRITON_CACHE_DIR=/home/user_lerobot/.cache/triton + +# Create the virtual environment +# We use a virtual environment inside the container—even though the container itself \ +# provides isolation—to closely resemble local development and allow users to \ +# run other Python projects in the same container without dependency conflicts. +RUN uv venv + +# Install Python dependencies for caching +COPY --chown=user_lerobot:user_lerobot pyproject.toml README.md MANIFEST.in ./ +# COPY --chown=user_lerobot:user_lerobot src/ src/ +RUN uv pip install --no-cache ".[smolvla]" + +RUN uv pip install rerun-sdk==0.24.0 scipy urdf_parser_py +ADD ./docker/visualize /workspace/visualize +CMD ["/bin/bash"] \ No newline at end of file diff --git a/docker/add_depth.py b/docker/add_depth.py new file mode 100644 index 0000000..2955f51 --- /dev/null +++ b/docker/add_depth.py @@ -0,0 +1,109 @@ +from lerobot.datasets.lerobot_dataset import LeRobotDataset +import numpy as np + +target_path = "docker/inputs/pick_orange_rs_depth" +src_ds = [ + ("datasets/realsense_output/pick_orange_0", "Pick the orange to the plate"), + ("datasets/realsense_output/pick_orange_1", "Pick the orange to the plate"), + ("datasets/realsense_output/pick_orange_2", "Pick the orange to the plate"), + ("datasets/realsense_output/pick_orange_3", "Pick the orange to the plate"), + ("datasets/realsense_output/pick_orange_4", "Pick the orange to the plate"), +] + +SINGLE_ARM_FEATURES = { + "action": { + "dtype": "float32", + "shape": (6,), + "names": [ + "shoulder_pan.pos", + "shoulder_lift.pos", + "elbow_flex.pos", + "wrist_flex.pos", + "wrist_roll.pos", + "gripper.pos", + ], + }, + "observation.state": { + "dtype": "float32", + "shape": (6,), + "names": [ + "shoulder_pan.pos", + "shoulder_lift.pos", + "elbow_flex.pos", + "wrist_flex.pos", + "wrist_roll.pos", + "gripper.pos", + ], + }, + "observation.images.front": { + "dtype": "image", + "shape": [480, 640, 3], + "names": ["height", "width", "channels"], + }, + "observation.images.front_depth": { + "dtype": "image", + "shape": [480, 640, 3], + "names": ["height", "width", "channels"], + }, +} + +from cloud_helper import Client +client = Client(host="localhost", port=50000) +# client = Client(host="120.48.81.132", port=50000) + +import os.path as osp +from os import system +from tqdm import tqdm + +if osp.exists(target_path): + system(f"rm -rf {target_path}") + +target = LeRobotDataset.create( + repo_id="O24H/Target", + fps=30, + root=target_path, + robot_type="so101_follower", + features=SINGLE_ARM_FEATURES, +) + + +for src_path, task in src_ds: + src = LeRobotDataset( + repo_id="O24H/Src", + root=src_path, + ) + + for eps_idx in range(src.num_episodes): + + frame_idx = range( + src.episode_data_index["from"][eps_idx].item(), + src.episode_data_index["to"][eps_idx].item(), + ) + + diff_actions = [src.__getitem__(i)["action"] - src.__getitem__(i - 1)["action"] for i in frame_idx if i > 0] + EPS = 1e-3 + + keep_idx = [i for i, a in zip(frame_idx, diff_actions) if (a.abs() > EPS).any()] + + compress_ratio = len(keep_idx) / len(frame_idx) + print(f"Episode {eps_idx}: compress ratio {compress_ratio:.2f}") + + if len(keep_idx) == 0: + continue + + for frame_idx in tqdm(keep_idx): + batch = src.__getitem__(frame_idx) + + front_img = batch["observation.images.front"].permute(1, 2, 0).contiguous().numpy() + front_depth = client.call_endpoint("get_depth", front_img) + + front_depth_ = front_depth[:, :, np.newaxis] * np.ones(3) * 255.0 + front_depth_ = front_depth_.astype(np.uint8) + + frame = { + "action": batch["action"], + "observation.state": batch["observation.state"], + "observation.images.front": front_img, + "observation.images.front_depth": front_depth_} + target.add_frame(frame, task=task) + target.save_episode() diff --git a/docker/cloud_helper.py b/docker/cloud_helper.py index 6b425ee..e68e6d6 100644 --- a/docker/cloud_helper.py +++ b/docker/cloud_helper.py @@ -2,14 +2,17 @@ import zmq import msgpack import msgpack_numpy as m -from typing import Any, Callable import logging +import time + +from typing import Any, Callable +import zstandard as zstd + logger = logging.getLogger(__name__) -import zstandard as zstd -compresser = zstd.ZstdCompressor(level=3) +compresser = zstd.ZstdCompressor(level=12) decompresser = zstd.ZstdDecompressor() @@ -102,6 +105,19 @@ class Client: raise Exception(f"Error from server: {message.get('data')}") + +def time_it(func): + def wrapper(*args, **kwargs): + start_time = time.time() + result = func(*args, **kwargs) + end_time = time.time() + elapsed_time = end_time - start_time + logger.info(f"'{func.__name__}' tooks {elapsed_time * 1000:.2f} ms") + return result + return wrapper + + + if __name__ == "__main__": import sys from time import sleep diff --git a/docker/infer.py b/docker/infer.py new file mode 100644 index 0000000..f8b6329 --- /dev/null +++ b/docker/infer.py @@ -0,0 +1,104 @@ +import json +import os +import signal +import threading +import time + +import torch + +from cloud_helper import Server +from lerobot.policies.factory import get_policy_class + +os.environ["CUDA_VISIBLE_DEVICES"] = "0" +os.environ["HF_HUB_OFFLINE"] = "1" +os.environ["HF_DATASETS_OFFLINE"] = "1" +# os.environ["HTTPS_PROXY"] = "http://192.168.16.68:18000" + +with open("/workspace/embolab/params/build_task.json") as f: + task_configs = json.load(f) + +model = task_configs["online_infer"]["model"] +assert model in ["smolvla", "act"], f"Unsupported model: {model}" + +checkpoint = task_configs["online_infer"]["checkpoint_path"] +if not checkpoint.endswith("/pretrained_model"): + checkpoint += "/pretrained_model" + print(f"Adjusted checkpoint path to: {checkpoint}") + +server_port = task_configs["online_infer"].get("port", 8080) + +class LerobotInferenceServer: + def __init__( + self, + checkpoint: str, + policy_type: str = "smolvla", + host: str = "localhost", + port: int = 5555, + device="cuda", + timeout: int = 3600, + ): + self.server = Server(host, port) + self.policy_type = policy_type + policy_class = get_policy_class(self.policy_type) + self.policy = policy_class.from_pretrained(checkpoint) + self.device = device + self.policy.to(self.device) + print(f"Loaded {self.policy_type.upper()} policy from {checkpoint}") + + self.timeout = timeout + self.last_activity = time.time() + self.stop_event = threading.Event() + self.monitor_thread = threading.Thread(target=self.watchout, daemon=True) + self.monitor_thread.start() + + def watchout(self): + while not self.stop_event.is_set(): + time.sleep(6) # Check every 6 seconds + elapsed = time.time() - self.last_activity + if elapsed > self.timeout: + print(f"No activity for {elapsed:.0f} seconds. Shutting down due to timeout.") + # Force exit since loop_forever might block + os.kill(os.getpid(), signal.SIGINT) + + def get_actions(self, batch): + # batch = { + # "observation": { + # "state": ..., + # "images.front": ..., HWC uint8 + # "images.wrist": ..., + # }, + # "instruction": ..., + # } + + obs = {} + + for k, v in batch["observation"].items(): + if k.startswith("images.") and v is not None: + img = v.astype("float32") / 255.0 + img = img.transpose(2, 0, 1) # HWC -> CHW + img = torch.from_numpy(img).unsqueeze(0).to(self.device) + obs[f"observation.{k}"] = img + elif k == "state": + tensor = torch.from_numpy(v.astype("float32")).unsqueeze(0).to(self.device) + obs[f"observation.{k}"] = tensor + obs["task"] = batch["instruction"] + + action_chunk = self.policy.predict_action_chunk(obs) + + self.last_activity = time.time() + + return action_chunk.cpu().numpy() # (B, chunk_size, action_dim) + + def run(self): + self.server.register_endpoint("get_actions", self.get_actions) + print(f"Lerobot {self.policy_type.upper()} Server is running...") + + self.server.loop_forever() + + +if __name__ == "__main__": + + server = LerobotInferenceServer( + checkpoint=checkpoint, policy_type=model, host="0.0.0.0", port=server_port, timeout=3600 + ) + server.run() diff --git a/docker/infer_task.json b/docker/infer_task.json new file mode 100644 index 0000000..5fa0e7e --- /dev/null +++ b/docker/infer_task.json @@ -0,0 +1,8 @@ +{ + "task_id": "b5c75014c1142feab3ee395b4a0bcc0", + "online_infer": { + "model": "smolvla", + "checkpoint_path": "/workspace/embolab/outputs/pick_toy_compact/checkpoints/020000", + "port": 8080 + } +} \ No newline at end of file diff --git a/docker/marigold_server.py b/docker/marigold_server.py new file mode 100644 index 0000000..4f1f093 --- /dev/null +++ b/docker/marigold_server.py @@ -0,0 +1,26 @@ +import os + +from diffusers import MarigoldDepthPipeline +import torch + +from cloud_helper import Server, time_it + +device = torch.device("cuda:1") +model = MarigoldDepthPipeline.from_pretrained("prs-eth/marigold-depth-v1-1").to(device) + +@time_it +def get_depth(image): + output = model(image) + depth = output.prediction.squeeze() + return depth + +import logging +logger = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) + +server = Server(host="0.0.0.0", port=50001) +server.register_endpoint("get_depth", get_depth) +print(f"Marigold Depth Pipeline is running at {server.host}:{server.port}...") +server.loop_forever() \ No newline at end of file diff --git a/docker/merge.py b/docker/merge.py index be619ea..ae5d61d 100644 --- a/docker/merge.py +++ b/docker/merge.py @@ -1,4 +1,5 @@ import json +import os from pathlib import Path from tqdm import tqdm @@ -6,12 +7,12 @@ from tqdm import tqdm from lerobot.datasets.lerobot_dataset import LeRobotDataset -with open("/workspace/inputs/task.json", "r") as f: +with open("/workspace/embolab/params/build_task.json", "r") as f: task_config = json.load(f) -src_dataset_paths = [i for i in Path(task_config["train"]["input_data_path"]).iterdir() if i.is_dir()] +src_dataset_paths = [i for i in Path(task_config["merge"]["input_data_path"]).iterdir() if i.is_dir()] -EPS = 1e-2 +# EPS = 1e-2 # Feature Check features = {} @@ -30,13 +31,14 @@ for p in src_dataset_paths: # assert dataset.features[k] == features[k], f"Feature key {k} mismatch in dataset {p}" # Initialize Target Dataset -target_path = Path(task_config["train"]["output_data_path"]) +target_path = Path(task_config["merge"]["output_data_path"]) # assert not target_path.exists(), f"Output path {target_path} already exists!" if target_path.exists(): import os - os.system(f"rm -rf {target_path}") +tmp_target_path = "/tmp/lerobot_merge_temp" + ### using images to store all data rather than videos: ### 35s per episode -> 20s per episode but size will be ~40x larger 6M -> 260M # for i in features.keys(): @@ -62,15 +64,18 @@ if target_path.exists(): target = LeRobotDataset.create( repo_id="O24H/Target", fps=30, - root=target_path, + root=tmp_target_path, robot_type="so101_follower", features=features, - image_writer_processes=8, image_writer_threads=16, ) for p in src_dataset_paths: - src = LeRobotDataset(repo_id="O24H/Src", root=p) + + try: + src = LeRobotDataset(repo_id="O24H/Src", root=p) + except: + print("Error while Processing: ", p, ", Skip...") for eps_idx in tqdm(range(src.num_episodes), desc=f"Processing episode in {p.name}"): frame_idx = range( @@ -78,20 +83,22 @@ for p in src_dataset_paths: src.episode_data_index["to"][eps_idx].item(), ) - eps_data = [src.__getitem__(i) for i in frame_idx] + # eps_data = [src.__getitem__(i) for i in frame_idx] - diff_actions = [eps_data[i]["action"] - eps_data[i - 1]["action"] for i in range(1, len(eps_data))] - keep_idx = [i + 1 for i, a in enumerate(diff_actions) if (a.abs() > EPS).any()] + # diff_actions = [eps_data[i]["action"] - eps_data[i - 1]["action"] for i in range(1, len(eps_data))] + # keep_idx = [i + 1 for i, a in enumerate(diff_actions) if (a.abs() > EPS).any()] - compress_ratio = len(keep_idx) / len(frame_idx) - print(f"Episode {eps_idx}: compress ratio {compress_ratio:.2f}") + # compress_ratio = len(keep_idx) / len(frame_idx) + # print(f"Episode {eps_idx}: compress ratio {compress_ratio:.2f}") + + keep_idx = frame_idx # No Compression if len(keep_idx) < 32: continue # Skip too short episodes after compression for o in keep_idx: - batch = eps_data[o] + batch = src.__getitem__(o) image_keys = [k for k in batch.keys() if k.startswith("observation.images.")] @@ -105,4 +112,25 @@ for p in src_dataset_paths: target.add_frame(frame, task=batch["task"]) + # manually copy the video and skip encoding process + for key in target.meta.video_keys: + __src_video_path = src.root / src.meta.get_video_file_path(eps_idx, key) + __tgt_video_path = target.root / target.meta.get_video_file_path(target.num_episodes, key) + __tgt_video_path.parent.mkdir(parents=True, exist_ok=True) + os.system(f"cp {__src_video_path} {__tgt_video_path}") + target.save_episode() + + # Remove images if video exists + for key in target.meta.video_keys: + __img_dir = target._get_image_file_path( + episode_index=target.num_episodes - 1, image_key=key, frame_index=0 + ).parent + if __img_dir.exists(): + print(f"Removing image dir: {__img_dir}") + os.system(f"rm -rf {__img_dir}") + + +# move tmp target to final target +os.system(f"mv {tmp_target_path} {target_path}") +print(f"Merged dataset saved at {target_path}") diff --git a/docker/merge_task.json b/docker/merge_task.json index a75eea4..e645e47 100644 --- a/docker/merge_task.json +++ b/docker/merge_task.json @@ -1,7 +1,7 @@ { "task_id": "b5c75014c1142feab3ee395b4a0bcc0", - "train": { - "input_data_path": "/workspace/inputs/", - "output_data_path": "/workspace/outputs/pick_orange_mixed" + "merge": { + "input_data_path": "/workspace/embolab/inputs/", + "output_data_path": "/workspace/embolab/outputs/pick_orange_new" } } \ No newline at end of file diff --git a/docker/moge2_server.py b/docker/moge2_server.py new file mode 100644 index 0000000..87411f6 --- /dev/null +++ b/docker/moge2_server.py @@ -0,0 +1,31 @@ +import os +import torch + +from cloud_helper import Server, time_it + +from moge.model.v2 import MoGeModel # Let's try MoGe-2 + +os.environ["CUDA_VISIBLE_DEVICES"] = "0" + +device = torch.device("cuda") +# model = MoGeModel.from_pretrained("Ruicheng/moge-2-vits-normal").to(device) +model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(device) + +def get_depth(image, normalize=True): + image = torch.tensor(image / 255, dtype=torch.float32, device=device).permute(2, 0, 1) + output = model.infer(image) + depth = output["depth"].cpu().numpy() + if normalize: + depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8) + return depth + +import logging +logger = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) + +server = Server(host="0.0.0.0", port=50001) +server.register_endpoint("get_depth", get_depth) +print(f"MoGe-2 Server is running at {server.host}:{server.port}...") +server.loop_forever() diff --git a/docker/smolvla_dataloader.py b/docker/smolvla_dataloader.py new file mode 100644 index 0000000..23c6bd3 --- /dev/null +++ b/docker/smolvla_dataloader.py @@ -0,0 +1,61 @@ + +from tqdm import tqdm +from cloud_helper import Client +from lerobot.datasets.lerobot_dataset import LeRobotDataset + +import cv2 +import numpy as np + + +client = Client(host="127.0.0.1", port=50000) + +src = LeRobotDataset( + repo_id="Foo/Bar", + root="./docker/inputs/pick_orange_vggt", # Modify this path accordingly +) + +EPS = 1e-2 # Threshold for action change + +for eps_idx in tqdm(range(src.num_episodes), desc="Processing episodes: "): + + frame_idx = range( + src.episode_data_index["from"][eps_idx].item(), + src.episode_data_index["to"][eps_idx].item(), + ) + + # eps_data = [src.__getitem__(i) for i in frame_idx] + + # diff_actions = [eps_data[i]["action"] - eps_data[i - 1]["action"] for i in range(1, len(eps_data))] + # keep_idx = [i + 1 for i, a in enumerate(diff_actions) if (a.abs() > EPS).any()] + + # compress_ratio = len(keep_idx) / len(frame_idx) + # print(f"Episode {eps_idx}: compress ratio {compress_ratio:.2f}") + + keep_idx = frame_idx # Keep all frames + + for o in keep_idx: + + # raw_batch = eps_data[o] + raw_batch = src.__getitem__(o) + + state_np = raw_batch["observation.state"].numpy() + image_front_np = (raw_batch["observation.images.front"].permute(1, 2, 0).numpy() * 255).astype("uint8") + image_wrist_np = (raw_batch["observation.images.front_depth"].permute(1, 2, 0).numpy() * 255).astype("uint8") + + batch = { + "observation": { + "images.front": image_front_np, + "images.wrist": np.zeros_like(image_front_np), # Placeholder for wrist image + "state": state_np, + }, + "instruction": raw_batch["task"], + } + + action_chunk = client.call_endpoint("get_actions", batch) + + target_action = raw_batch["action"] + predicted_action = action_chunk[0] + + print(f"Target action: {target_action}, Predicted action: {predicted_action}") + + pass \ No newline at end of file diff --git a/docker/smolvla_executor.py b/docker/smolvla_executor.py index 5f27238..d515601 100644 --- a/docker/smolvla_executor.py +++ b/docker/smolvla_executor.py @@ -91,7 +91,7 @@ class SmolVLAExecutor: def loop_once(self): if len(self._action_queue) <= 1: new_actions = self.get_actions() - self._action_queue.extend(new_actions.transpose(0, 1)) + self._action_queue.extend(new_actions) # Apply the filter self.apply_filter() diff --git a/docker/smolvla_server.py b/docker/smolvla_server.py index 3480a40..cc11668 100644 --- a/docker/smolvla_server.py +++ b/docker/smolvla_server.py @@ -4,7 +4,7 @@ import os from cloud_helper import Server from lerobot.policies.factory import get_policy_class -os.environ["CUDA_VISIBLE_DEVICES"] = "0" +os.environ["CUDA_VISIBLE_DEVICES"] = "3" os.environ["HF_HUB_OFFLINE"] = "1" @@ -54,12 +54,12 @@ class LerobotInferenceServer: def run(self): self.server.register_endpoint("get_actions", self.get_actions) - print(f"Lerobot {self.policy_type.upper()} Server is running...") + print(f"Lerobot {self.policy_type.upper()} Server is running at {self.server.host}:{self.server.port}...") self.server.loop_forever() if __name__ == "__main__": - smolvla_checkpoint = "./20250901/pick_red_marker_smolvla/checkpoints/last/pretrained_model" + smolvla_checkpoint = "202501101/pick_toys_01/checkpoints/010000" + "/pretrained_model" server = LerobotInferenceServer( checkpoint=smolvla_checkpoint, policy_type="smolvla", host="0.0.0.0", port=50000 ) diff --git a/docker/train.py b/docker/train.py index b142619..24e80ff 100644 --- a/docker/train.py +++ b/docker/train.py @@ -1,7 +1,7 @@ import json import os -with open("/workspace/inputs/task.json") as f: +with open("/workspace/embolab/params/build_task.json") as f: task_configs = json.load(f) @@ -12,11 +12,13 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false" assert "train" in task_configs, "Not a validate train config" assert task_configs["train"]["model"] in ["act", "smolvla"], "Only act and smolvla are supported for training" -use_policy = ( - "--policy.path=lerobot/smolvla_base" - if task_configs["train"]["model"] == "smolvla" - else "--policy.type=act" -) +# use_policy = ( +# "--policy.path=lerobot/smolvla_base" +# if task_configs["train"]["model"] == "smolvla" +# else "--policy.type=act" +# ) +use_policy = f"--policy.type={task_configs['train']['model']}" + task_id = task_configs["task_id"] data_path = task_configs["train"]["input_data_path"] ckpt_path = task_configs["train"]["checkpoint_path"] @@ -50,4 +52,8 @@ train_cmd = f"""lerobot-train \ """ print("Executing command:\n", train_cmd) -os.system(train_cmd) + +import subprocess, sys + +completed = subprocess.run(train_cmd, shell=True) +sys.exit(completed.returncode) diff --git a/docker/train_task.json b/docker/train_task.json index 6db760f..ffae72a 100644 --- a/docker/train_task.json +++ b/docker/train_task.json @@ -3,10 +3,10 @@ "gpu_id": "0", "train": { "model":"smolvla", - "epochs":20, - "batch_size":64, - "log_path": "/workspace/logs", - "checkpoint_path": "/workspace/outputs/checkpoints", - "input_data_path": "/workspace/inputs/pick_red_marker" + "epochs":200, + "batch_size":256, + "log_path": "/workspace/embolab/logs", + "checkpoint_path": "/workspace/embolab/outputs/pick_orange_new", + "input_data_path": "/workspace/embolab/inputs/pick_orange_new" } } \ No newline at end of file diff --git a/docker/visualize/assets/SO101/README.md b/docker/visualize/assets/SO101/README.md new file mode 100644 index 0000000..44d954d --- /dev/null +++ b/docker/visualize/assets/SO101/README.md @@ -0,0 +1,35 @@ +# SO101 Robot - URDF and MuJoCo Description + +This repository contains the URDF and MuJoCo (MJCF) files for the SO101 robot. + +## Overview + +- The robot model files were generated using the [onshape-to-robot](https://github.com/Rhoban/onshape-to-robot) plugin from a CAD model designed in Onshape. +- The generated URDFs were modified to allow meshes with relative paths instead of `package://...`. +- Base collision meshes were removed due to problematic collision behavior during simulation and planning. + +## Calibration Methods + +The MuJoCo file `scene.xml` supports two differenly calibrated SO101 robot files: + +- **New Calibration (Default)**: Each joint's virtual zero is set to the **middle** of its joint range. Use -> `so101_new_calib.xml`. +- **Old Calibration**: Each joint's virtual zero is set to the configuration where the robot is **fully extended horizontally**. Use -> `so101_old_calib.xml`. + +To switch between calibration methods, modify the included robot file in `scene.xml`. + +## Motor Parameters + +Motor properties for the STS3215 motors used in the robot are adapted from the [Open Duck Mini project](https://github.com/apirrone/Open_Duck_Mini). + +## Gripper Note + +In LeRobot, the gripper is represented as a **linear joint**, where: + +* `0` = fully closed +* `100` = fully open + +This mapping is **not yet reflected** in the current URDF and MuJoCo files. + +--- + +Feel free to open an issue or contribute improvements! diff --git a/docker/visualize/assets/SO101/assets/base_motor_holder_so101_v1.part b/docker/visualize/assets/SO101/assets/base_motor_holder_so101_v1.part new file mode 100644 index 0000000..9ea635f --- /dev/null +++ b/docker/visualize/assets/SO101/assets/base_motor_holder_so101_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "84d8ae1881704ebae1ffb70a", + "documentMicroversion": "0eea3500852bdb2f58b1cb79", + "documentVersion": "a5c3b0dfaa52ddd6829011cd", + "elementId": "22efbe4e0bef24fcd20f96e5", + "fullConfiguration": "default", + "id": "MCOhripg0ry51VlsC", + "isStandardContent": false, + "name": "Base_motor_holder_SO101 v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/base_motor_holder_so101_v1.stl b/docker/visualize/assets/SO101/assets/base_motor_holder_so101_v1.stl new file mode 100644 index 0000000..ac9c380 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/base_motor_holder_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8cd2f241037ea377af1191fffe0dd9d9006beea6dcc48543660ed41647072424 +size 1877084 diff --git a/docker/visualize/assets/SO101/assets/base_so101_v2.part b/docker/visualize/assets/SO101/assets/base_so101_v2.part new file mode 100644 index 0000000..39b1553 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/base_so101_v2.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "bf61a6bc85b1d1a8bf9ea51b", + "documentMicroversion": "20484d37162a32a8a41a37f2", + "documentVersion": "25801b070e5b360715de8a30", + "elementId": "312f32f0073fa6e8e36fba7a", + "fullConfiguration": "default", + "id": "MY69cJlqvSzIiODdH", + "isStandardContent": false, + "name": "Base_SO101 v2 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/base_so101_v2.stl b/docker/visualize/assets/SO101/assets/base_so101_v2.stl new file mode 100644 index 0000000..503d30b --- /dev/null +++ b/docker/visualize/assets/SO101/assets/base_so101_v2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb12b7026575e1f70ccc7240051f9d943553bf34e5128537de6cd86fae33924d +size 471584 diff --git a/docker/visualize/assets/SO101/assets/motor_holder_so101_base_v1.part b/docker/visualize/assets/SO101/assets/motor_holder_so101_base_v1.part new file mode 100644 index 0000000..74dfba9 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/motor_holder_so101_base_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "652d5731024e57367badfda6", + "documentMicroversion": "56a8b8013480c176fd87df8d", + "documentVersion": "984ac31c92cac3664c8effb3", + "elementId": "6fb7b7f9315511b548d670ff", + "fullConfiguration": "default", + "id": "Mf4ZebMr4BkShucFj", + "isStandardContent": false, + "name": "Motor_holder_SO101_Base v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/motor_holder_so101_base_v1.stl b/docker/visualize/assets/SO101/assets/motor_holder_so101_base_v1.stl new file mode 100644 index 0000000..f8e3d75 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/motor_holder_so101_base_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:31242ae6fb59d8b15c66617b88ad8e9bded62d57c35d11c0c43a70d2f4caa95b +size 1129384 diff --git a/docker/visualize/assets/SO101/assets/motor_holder_so101_wrist_v1.part b/docker/visualize/assets/SO101/assets/motor_holder_so101_wrist_v1.part new file mode 100644 index 0000000..7df4b34 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/motor_holder_so101_wrist_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "4bd66da73cacb4d946d43e44", + "documentMicroversion": "2bf56247e58b70e90806e318", + "documentVersion": "df78bb7089f1de7d5588d238", + "elementId": "d7dfe76e402c21bbd8124e43", + "fullConfiguration": "default", + "id": "MN9BZ1p69dQQtKTjq", + "isStandardContent": false, + "name": "Motor_holder_SO101_Wrist v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/motor_holder_so101_wrist_v1.stl b/docker/visualize/assets/SO101/assets/motor_holder_so101_wrist_v1.stl new file mode 100644 index 0000000..e55b719 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/motor_holder_so101_wrist_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:887f92e6013cb64ea3a1ab8675e92da1e0beacfd5e001f972523540545e08011 +size 1052184 diff --git a/docker/visualize/assets/SO101/assets/moving_jaw_so101_v1.part b/docker/visualize/assets/SO101/assets/moving_jaw_so101_v1.part new file mode 100644 index 0000000..a294510 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/moving_jaw_so101_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "46218c02ef80d36172edbb35", + "documentMicroversion": "68b7d387e2500c451586ae59", + "documentVersion": "79c101d1a0207b77362b561a", + "elementId": "d4b1411d5d7333298f6e2458", + "fullConfiguration": "default", + "id": "MrHPLr9hZkrXwcSA4", + "isStandardContent": false, + "name": "Moving_Jaw_SO101 v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/moving_jaw_so101_v1.stl b/docker/visualize/assets/SO101/assets/moving_jaw_so101_v1.stl new file mode 100644 index 0000000..eb17d25 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/moving_jaw_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:785a9dded2f474bc1d869e0d3dae398a3dcd9c0c345640040472210d2861fa9d +size 1413584 diff --git a/docker/visualize/assets/SO101/assets/rotation_pitch_so101_v1.part b/docker/visualize/assets/SO101/assets/rotation_pitch_so101_v1.part new file mode 100644 index 0000000..b873134 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/rotation_pitch_so101_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "14078aa6723c502d07d6902e", + "documentMicroversion": "c0fca717407275159bcc6ed7", + "documentVersion": "3d9a887ff68fa477d98162b8", + "elementId": "43d24b3857ff686b275578bf", + "fullConfiguration": "default", + "id": "MrQ6Kmk9QDZlwbp95", + "isStandardContent": false, + "name": "Rotation_Pitch_SO101 v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/rotation_pitch_so101_v1.stl b/docker/visualize/assets/SO101/assets/rotation_pitch_so101_v1.stl new file mode 100644 index 0000000..b536cb4 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/rotation_pitch_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9be900cc2a2bf718102841ef82ef8d2873842427648092c8ed2ca1e2ef4ffa34 +size 883684 diff --git a/docker/visualize/assets/SO101/assets/sts3215_03a_no_horn_v1.part b/docker/visualize/assets/SO101/assets/sts3215_03a_no_horn_v1.part new file mode 100644 index 0000000..f5483ea --- /dev/null +++ b/docker/visualize/assets/SO101/assets/sts3215_03a_no_horn_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "56e5f3702dad85e17841d2e2", + "documentMicroversion": "7958a6acbc8e0d0a0a611746", + "documentVersion": "29a4c51b8bf277a22743a333", + "elementId": "8c14fb13a6557ec89ff5d227", + "fullConfiguration": "default", + "id": "MOcaIFg8XgL+Ybg9z", + "isStandardContent": false, + "name": "STS3215_03a_no_horn v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/sts3215_03a_no_horn_v1.stl b/docker/visualize/assets/SO101/assets/sts3215_03a_no_horn_v1.stl new file mode 100644 index 0000000..18e9335 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/sts3215_03a_no_horn_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75ef3781b752e4065891aea855e34dc161a38a549549cd0970cedd07eae6f887 +size 865884 diff --git a/docker/visualize/assets/SO101/assets/sts3215_03a_v1.part b/docker/visualize/assets/SO101/assets/sts3215_03a_v1.part new file mode 100644 index 0000000..d5ff13e --- /dev/null +++ b/docker/visualize/assets/SO101/assets/sts3215_03a_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "d2941bdba816affebdc6d6f0", + "documentMicroversion": "5904ef3cea04a0d0bc88b698", + "documentVersion": "dd4f7470101215836a4ae8c9", + "elementId": "e670b72d49b06f88fad5dbd8", + "fullConfiguration": "default", + "id": "M5vQNpe0onRFueych", + "isStandardContent": false, + "name": "STS3215_03a v1 <5>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/sts3215_03a_v1.stl b/docker/visualize/assets/SO101/assets/sts3215_03a_v1.stl new file mode 100644 index 0000000..a14c57b --- /dev/null +++ b/docker/visualize/assets/SO101/assets/sts3215_03a_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a37c871fb502483ab96c256baf457d36f2e97afc9205313d9c5ab275ef941cd0 +size 954084 diff --git a/docker/visualize/assets/SO101/assets/under_arm_so101_v1.part b/docker/visualize/assets/SO101/assets/under_arm_so101_v1.part new file mode 100644 index 0000000..bef615a --- /dev/null +++ b/docker/visualize/assets/SO101/assets/under_arm_so101_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "9f5d6db47eb112442b9f130f", + "documentMicroversion": "e99cf45162e34789bd99512b", + "documentVersion": "817ebf29c5663d412edc0753", + "elementId": "2813aaffe3c8a342616d3527", + "fullConfiguration": "default", + "id": "M9yAEiX02J3c4HqXa", + "isStandardContent": false, + "name": "Under_arm_SO101 v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/under_arm_so101_v1.stl b/docker/visualize/assets/SO101/assets/under_arm_so101_v1.stl new file mode 100644 index 0000000..47b611e --- /dev/null +++ b/docker/visualize/assets/SO101/assets/under_arm_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d01d1f2de365651dcad9d6669e94ff87ff7652b5bb2d10752a66a456a86dbc71 +size 1975884 diff --git a/docker/visualize/assets/SO101/assets/upper_arm_so101_v1.part b/docker/visualize/assets/SO101/assets/upper_arm_so101_v1.part new file mode 100644 index 0000000..9198a2b --- /dev/null +++ b/docker/visualize/assets/SO101/assets/upper_arm_so101_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "57f3eae43434311c28ac752b", + "documentMicroversion": "33eeab136e831427f0f0ca74", + "documentVersion": "435d47b71ef26075bf82672c", + "elementId": "a8e0c02dc43f7ccb373c52e4", + "fullConfiguration": "default", + "id": "Ml3rwO4kV53jDRgcs", + "isStandardContent": false, + "name": "Upper_arm_SO101 v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/upper_arm_so101_v1.stl b/docker/visualize/assets/SO101/assets/upper_arm_so101_v1.stl new file mode 100644 index 0000000..8832740 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/upper_arm_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:475056e03a17e71919b82fd88ab9a0b898ab50164f2a7943652a6b2941bb2d4f +size 1303484 diff --git a/docker/visualize/assets/SO101/assets/waveshare_mounting_plate_so101_v2.part b/docker/visualize/assets/SO101/assets/waveshare_mounting_plate_so101_v2.part new file mode 100644 index 0000000..c71d889 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/waveshare_mounting_plate_so101_v2.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "066f8b5064455ec46759cd8c", + "documentMicroversion": "04c5790374bf3edfbbb7e818", + "documentVersion": "408440a116f7d8700bbb11c2", + "elementId": "dc35e56269e36de39738b34d", + "fullConfiguration": "default", + "id": "MjhXxhyF1+iAgCtUh", + "isStandardContent": false, + "name": "WaveShare_Mounting_Plate_SO101 v2 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/waveshare_mounting_plate_so101_v2.stl b/docker/visualize/assets/SO101/assets/waveshare_mounting_plate_so101_v2.stl new file mode 100644 index 0000000..e0d90d5 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/waveshare_mounting_plate_so101_v2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e197e24005a07d01bbc06a8c42311664eaeda415bf859f68fa247884d0f1a6e9 +size 62784 diff --git a/docker/visualize/assets/SO101/assets/wrist_roll_follower_so101_v1.part b/docker/visualize/assets/SO101/assets/wrist_roll_follower_so101_v1.part new file mode 100644 index 0000000..40d1cf5 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/wrist_roll_follower_so101_v1.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "e02f1e1d3fdd766a19a55890", + "documentMicroversion": "03f1dfc090db6bbecdb14475", + "documentVersion": "8a15327cfbde0344e0951076", + "elementId": "2317bd70c68862eeebd64492", + "fullConfiguration": "default", + "id": "MpI0voU28BOAZ6D9x", + "isStandardContent": false, + "name": "Wrist_Roll_Follower_SO101 v1 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/wrist_roll_follower_so101_v1.stl b/docker/visualize/assets/SO101/assets/wrist_roll_follower_so101_v1.stl new file mode 100644 index 0000000..9a5fa8f --- /dev/null +++ b/docker/visualize/assets/SO101/assets/wrist_roll_follower_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4b17b410a12d64ec39554abc3e8054d8a97384b2dc4a8d95a5ecb2a93670f5f4 +size 1439884 diff --git a/docker/visualize/assets/SO101/assets/wrist_roll_pitch_so101_v2.part b/docker/visualize/assets/SO101/assets/wrist_roll_pitch_so101_v2.part new file mode 100644 index 0000000..1c099e9 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/wrist_roll_pitch_so101_v2.part @@ -0,0 +1,14 @@ +{ + "configuration": "default", + "documentId": "eb144d215e733b8dbbb50b81", + "documentMicroversion": "4fef760722dee3a9b5ff19b1", + "documentVersion": "5880c1e9413206cac10772d0", + "elementId": "3c22c2c23cb0ce545b9df2ba", + "fullConfiguration": "default", + "id": "Ma99J59HxnSe2TArb", + "isStandardContent": false, + "name": "Wrist_Roll_Pitch_SO101 v2 <1>", + "partId": "JFD", + "suppressed": false, + "type": "Part" +} \ No newline at end of file diff --git a/docker/visualize/assets/SO101/assets/wrist_roll_pitch_so101_v2.stl b/docker/visualize/assets/SO101/assets/wrist_roll_pitch_so101_v2.stl new file mode 100644 index 0000000..2f53171 --- /dev/null +++ b/docker/visualize/assets/SO101/assets/wrist_roll_pitch_so101_v2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c7ec5525b4d8b9e397a30ab4bb0037156a5d5f38a4adf2c7d943d6c56eda5ae +size 2699784 diff --git a/docker/visualize/assets/SO101/joints_properties.xml b/docker/visualize/assets/SO101/joints_properties.xml new file mode 100644 index 0000000..0e2b2a5 --- /dev/null +++ b/docker/visualize/assets/SO101/joints_properties.xml @@ -0,0 +1,12 @@ + + + + + + + + + + + \ No newline at end of file diff --git a/docker/visualize/assets/SO101/scene.xml b/docker/visualize/assets/SO101/scene.xml new file mode 100644 index 0000000..255b970 --- /dev/null +++ b/docker/visualize/assets/SO101/scene.xml @@ -0,0 +1,24 @@ + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/docker/visualize/assets/SO101/so101_new_calib.urdf b/docker/visualize/assets/SO101/so101_new_calib.urdf new file mode 100644 index 0000000..9552a23 --- /dev/null +++ b/docker/visualize/assets/SO101/so101_new_calib.urdf @@ -0,0 +1,453 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + \ No newline at end of file diff --git a/docker/visualize/assets/SO101/so101_new_calib.xml b/docker/visualize/assets/SO101/so101_new_calib.xml new file mode 100644 index 0000000..e24ff77 --- /dev/null +++ b/docker/visualize/assets/SO101/so101_new_calib.xml @@ -0,0 +1,162 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docker/visualize/assets/SO101/so101_old_calib.urdf b/docker/visualize/assets/SO101/so101_old_calib.urdf new file mode 100644 index 0000000..c75bb38 --- /dev/null +++ b/docker/visualize/assets/SO101/so101_old_calib.urdf @@ -0,0 +1,435 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + diff --git a/docker/visualize/assets/SO101/so101_old_calib.xml b/docker/visualize/assets/SO101/so101_old_calib.xml new file mode 100644 index 0000000..7bc79e9 --- /dev/null +++ b/docker/visualize/assets/SO101/so101_old_calib.xml @@ -0,0 +1,160 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docker/visualize/visualize.py b/docker/visualize/visualize.py new file mode 100644 index 0000000..85f07f4 --- /dev/null +++ b/docker/visualize/visualize.py @@ -0,0 +1,108 @@ +from pathlib import Path + +import numpy as np +import rerun as rr +import rerun.blueprint as rrb +import scipy.spatial.transform as st + +from lerobot.datasets.lerobot_dataset import LeRobotDataset +from urdf_parser_py import urdf as urdf_parser + +SO101_URDF = Path(__file__).resolve().parent / "assets/SO101/so101_new_calib.urdf" +DEFAULT_JOINT_NAMES = [ + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", + "gripper", +] + + +def visualize_episode(dataset_path: str | Path, episode_index: int): + # Init rerun env + bp = rrb.Blueprint(collapse_panels=True) + rr.init("play_lerobot", spawn=False, default_blueprint=bp) + rr.send_blueprint(bp, make_active=True) + rr.serve_grpc(server_memory_limit="80%") + + # Visualize URDF + rr.log_file_from_path(file_path=SO101_URDF) + rr.set_time("frame", sequence=0) + + # URDF Controller + urdf = urdf_parser.URDF.from_xml_file(SO101_URDF) + + def get_joint_paths(urdf: urdf_parser.URDF, entity_prefix: str = "") -> dict[str, str]: + joint_paths = {} + + def get_joint_path(joint: urdf_parser.Joint) -> str: + """Return the entity path for the URDF joint.""" + root_name = urdf.get_root() + joint_names = urdf.get_chain(root_name, joint.child) + return entity_prefix + ("/".join(joint_names)) if entity_prefix else "/" + ("/".join(joint_names)) + + for joint in urdf.joints: + joint_paths[joint.name] = get_joint_path(joint) + return joint_paths + + joint_paths = get_joint_paths(urdf, entity_prefix="/so101_new_calib/") + + ds = LeRobotDataset(repo_id="O24H/Vis", root=dataset_path) + + frame_idx = range( + ds.episode_data_index["from"][episode_index].item(), + ds.episode_data_index["to"][episode_index].item(), + ) + + for i in frame_idx: + frame = ds[i] + + # dt = 1 / ds.fps + # rr.set_time_seconds("timestamp", i * dt) + + rr.set_time("frame", sequence=i + 1) + + # Visualize robot state + for joint_index, joint_name in enumerate(JOINT_NAMES): + joint = urdf.joint_map[joint_name] + if joint.type == "revolute": + # 取得关节可以旋转的轴 + axis = np.array(getattr(joint, "axis", [0, 0, 1])) # 默认轴为z轴 + + ang = frame["observation.state"][joint_index] + + R = rr.RotationAxisAngle(axis=axis, degrees=ang) + + entity_path = joint_paths[joint_name] + + rr.log( + entity_path, + rr.Transform3D(rotation=R), + ) + + for k in ds.meta.camera_keys: + rr.log( + "/camera/" + k, + rr.Image((frame[k] * 255).permute(1, 2, 0).numpy().astype(np.uint8)), + ) + + +if __name__ == "__main__": + import argparse + import time + + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, required=True, help="Path to dataset") + parser.add_argument("--episode", type=int, default=0, help="Episode to visualize") + args = parser.parse_args() + + tic = time.time() + + visualize_episode(dataset_path=args.dataset, episode_index=args.episode - 1) # args.episode start from 1 + + toc = time.time() + print(f"Visualization server started in {toc - tic:.2f} seconds.") + + while True: + time.sleep(3600) diff --git a/server.sh b/server.sh new file mode 100644 index 0000000..bc257be --- /dev/null +++ b/server.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +source /home/hegn/miniconda3/bin/activate lerobot + +export CUDA_VISIBLE_DEVICES=1 +export HF_HUB_OFFLINE=1 + +# Inference on server +python src/lerobot/scripts/server/policy_server.py \ + --host=0.0.0.0 \ + --port=28080 \ No newline at end of file diff --git a/train.sh b/train.sh new file mode 100755 index 0000000..139ab21 --- /dev/null +++ b/train.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +source /home/hegn/miniconda3/bin/activate lerobot + +export HF_HUB_OFFLINE=1 +export HF_DATASET_OFFLINE=1 +export TOKENIZERS_PARALLELISM=false + +CUDA_VISIBLE_DEVICES=7 lerobot-train \ + --policy.path=lerobot/smolvla_base \ + --policy.push_to_hub=false \ + --dataset.repo_id=Open24H/dummy_ds \ + --dataset.root=./datasets/pick_toy_compact \ + --dataset.image_transforms.enable=true \ + --batch_size=128 \ + --output_dir=./202501101/pick_toy_compact \ + --steps=20000 --save_freq=2000 2>&1 | tee ./202501101/pick_toy_compact.log \ No newline at end of file