67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
import torch
|
|
import os
|
|
|
|
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"
|
|
|
|
|
|
class LerobotInferenceServer:
|
|
def __init__(
|
|
self,
|
|
checkpoint: str,
|
|
policy_type: str = "smolvla",
|
|
host: str = "localhost",
|
|
port: int = 5555,
|
|
device="cuda",
|
|
):
|
|
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}")
|
|
|
|
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)
|
|
|
|
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__":
|
|
smolvla_checkpoint = "./20250901/pick_red_marker_smolvla/checkpoints/last/pretrained_model"
|
|
server = LerobotInferenceServer(
|
|
checkpoint=smolvla_checkpoint, policy_type="smolvla", host="0.0.0.0", port=50000
|
|
)
|
|
server.run()
|