![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250529123619-683854c3c804e.png\")
<\/p>\n\u00a0\u03c0\u2080 \u4e3b\u63a7\u811a\u672c\u90fd\u5728 scripts \u4e2d：<\/p>\n
<\/p>\n
\u5176\u4e2d，serve_policy.py \u662f openpi \u4e2d\u7684\u7b56\u7565\u63a8\u7406\u670d\u52a1\u7aef\u811a\u672c，\u4f5c\u7528\u4e3a：\u542f\u52a8\u4e00\u4e2a WebSocket \u670d\u52a1\u5668，\u52a0\u8f7d\u9884\u8bad\u7ec3\u7b56\u7565\u6a21\u578b，\u7b49\u5f85\u5916\u90e8\u8bf7\u6c42（\u5982\u6765\u81ea main.py \u7684\u63a7\u5236\u7a0b\u5e8f），\u7136\u540e\u6267\u884c\u52a8\u4f5c\u63a8\u7406\u5e76\u8fd4\u56de\u7ed3\u679c<\/p>\n
\u4e00\u53e5\u8bdd\u603b\u7ed3\u4e00\u4e0b：<\/p>\n
\u5c06\u4e00\u4e2a Pi0 \u7b56\u7565\u6a21\u578b\u90e8\u7f72\u4e3a\u7f51\u7edc\u670d\u52a1（WebSocket API），\u4f9b\u673a\u5668\u4eba\u4e3b\u63a7\u8fdb\u7a0b\u8fdc\u7a0b\u8c03\u7528<\/p>\n<\/p>\n
\u76ee\u5f55<\/p>\n
1 \u5e93\u5f15\u7528<\/p>\n
2\u00a0\u53c2\u6570\u5b9a\u4e49<\/p>\n
3 \u52a0\u8f7d\u7b56\u7565\u6a21\u578b<\/p>\n
4\u00a0\u542f\u52a8\u63a8\u7406\u670d\u52a1<\/p>\n
5\u00a0\u4f7f\u7528\u65b9\u6cd5\u603b\u7ed3<\/p>\n
5.1 \u4f7f\u7528\u65b9\u6cd5<\/p>\n
5.2 \u95ee\u9898\u5206\u6790<\/p>\n
import dataclasses # \u7528\u4e8e\u521b\u5efa\u7ed3\u6784\u5316\u53c2\u6570\u5bf9\u8c61
\nimport enum # \u7528\u4e8e\u5b9a\u4e49\u73af\u5883\u679a\u4e3e\u7c7b\u578b
\nimport logging # \u7528\u4e8e\u65e5\u5fd7\u8f93\u51fa
\nimport socket # \u83b7\u53d6\u672c\u673a IP \u5730\u5740\u4fe1\u606f<\/p>\n
import tyro # \u547d\u4ee4\u884c\u53c2\u6570\u89e3\u6790\u5e93，\u6bd4 argparse \u66f4\u73b0\u4ee3<\/p>\n
# \u5f15\u5165\u7b56\u7565\u548c\u7b56\u7565\u914d\u7f6e\u6a21\u5757
\nfrom openpi.policies import policy as _policy
\nfrom openpi.policies import policy_config as _policy_config
\n# \u5f15\u5165 WebSocket \u63a8\u7406\u670d\u52a1\u6a21\u5757
\nfrom openpi.serving import websocket_policy_server
\n# \u5f15\u5165\u8bad\u7ec3\u914d\u7f6e\u6a21\u5757
\nfrom openpi.training import config as _config <\/p>\n
# \u5b9a\u4e49\u652f\u6301\u7684\u673a\u5668\u4eba\u73af\u5883\u679a\u4e3e\u7c7b\u578b
\nclass EnvMode(enum.Enum):
\n ALOHA = "aloha"
\n ALOHA_SIM = "aloha_sim"
\n DROID = "droid"
\n LIBERO = "libero"<\/p>\n
# \u5b9a\u4e49\u7528\u4e8e\u52a0\u8f7d checkpoint \u7684\u53c2\u6570\u7ed3\u6784
\n@dataclasses.dataclass
\nclass Checkpoint:
\n config: str # \u6a21\u578b\u5bf9\u5e94\u7684\u8bad\u7ec3\u914d\u7f6e\u540d，\u5982 "pi0_aloha_sim"
\n dir: str # checkpoint \u76ee\u5f55\u8def\u5f84（\u53ef\u4ee5\u662f\u672c\u5730\u6216 S3）<\/p>\n
# \u5b9a\u4e49\u9ed8\u8ba4\u7b56\u7565\u7c7b\u578b\u5360\u4f4d\u7b26\u7ed3\u6784
\n@dataclasses.dataclass
\nclass Default:
\n pass # \u4f7f\u7528\u9ed8\u8ba4\u6a21\u578b\u914d\u7f6e（\u4f8b\u5982\u4ece DEFAULT_CHECKPOINT \u4e2d\u9009\u62e9）<\/p>\n
# \u5b9a\u4e49\u4e3b\u5165\u53e3\u53c2\u6570\u7ed3\u6784\u4f53
\n@dataclasses.dataclass
\nclass Args:
\n env: EnvMode = EnvMode.ALOHA_SIM # \u4f7f\u7528\u7684\u673a\u5668\u4eba\u73af\u5883
\n default_prompt: str | None = None # \u9ed8\u8ba4\u8bed\u8a00\u63d0\u793a\u8bcd
\n port: int = 8000 # WebSocket \u670d\u52a1\u7aef\u53e3
\n record: bool = False # \u662f\u5426\u8bb0\u5f55\u7b56\u7565\u8f93\u51fa（\u7528\u4e8e\u8c03\u8bd5）
\n policy: Checkpoint | Default = dataclasses.field(default_factory=Default)<\/p>\n
# \u5b9a\u4e49\u6bcf\u4e2a\u73af\u5883\u5bf9\u5e94\u7684\u9ed8\u8ba4 checkpoint
\nDEFAULT_CHECKPOINT: dict[EnvMode, Checkpoint] = {
\n EnvMode.ALOHA: Checkpoint(
\n config="pi0_aloha",
\n dir="s3:\/\/openpi-assets\/checkpoints\/pi0_base",
\n ),
\n EnvMode.ALOHA_SIM: Checkpoint(
\n config="pi0_aloha_sim",
\n dir="s3:\/\/openpi-assets\/checkpoints\/pi0_aloha_sim",
\n ),
\n EnvMode.DROID: Checkpoint(
\n config="pi0_fast_droid",
\n dir="s3:\/\/openpi-assets\/checkpoints\/pi0_fast_droid",
\n ),
\n EnvMode.LIBERO: Checkpoint(
\n config="pi0_fast_libero",
\n dir="s3:\/\/openpi-assets\/checkpoints\/pi0_fast_libero",
\n ),
\n} <\/p>\n
\u5176\u4e2d，class Args\u00a0\u4f7f\u7528 Tyro \u6765\u4ece\u547d\u4ee4\u884c\u89e3\u6790\u53c2\u6570，\u5305\u62ec\u4ee5\u4e0b\u9009\u9879：<\/p>\n
# \u4ece\u9ed8\u8ba4\u73af\u5883\u521b\u5efa\u7b56\u7565\u6a21\u578b
\ndef create_default_policy(env: EnvMode, *, default_prompt: str | None = None) -> _policy.Policy:
\n if checkpoint := DEFAULT_CHECKPOINT.get(env):
\n return _policy_config.create_trained_policy(
\n _config.get_config(checkpoint.config), # \u52a0\u8f7d\u914d\u7f6e
\n checkpoint.dir, # \u52a0\u8f7d\u53c2\u6570
\n default_prompt=default_prompt # \u8bbe\u7f6e\u9ed8\u8ba4\u8bed\u8a00\u63d0\u793a
\n )
\n raise ValueError(f"Unsupported environment mode: {env}")<\/p>\n
# \u4ece\u547d\u4ee4\u884c\u53c2\u6570\u521b\u5efa\u7b56\u7565\u6a21\u578b
\ndef create_policy(args: Args) -> _policy.Policy:
\n match args.policy:
\n case Checkpoint(): # \u5982\u679c\u663e\u5f0f\u6307\u5b9a checkpoint
\n return _policy_config.create_trained_policy(
\n _config.get_config(args.policy.config),
\n args.policy.dir,
\n default_prompt=args.default_prompt
\n )
\n case Default(): # \u5426\u5219\u4f7f\u7528\u9ed8\u8ba4\u6a21\u578b
\n return create_default_policy(args.env, default_prompt=args.default_prompt) <\/p>\n
\u5206\u4e24\u79cd\u60c5\u51b5：<\/p>\n
1. \u624b\u52a8\u6307\u5b9a\u6a21\u578b\u4e0e\u8def\u5f84，\u5373\u589e\u52a0\u53c2\u6570 –policy：<\/p>\n
–policy 'Checkpoint(config="pi0_aloha", dir="s3:\/\/openpi-assets\/checkpoints\/pi0_base")' <\/p>\n
2. \u4f7f\u7528\u5185\u7f6e\u9ed8\u8ba4\u7b56\u7565（\u4f1a\u4ece S3 \u4e0b\u8f7d），\u5373\u4e0d\u6307\u5b9a\u53c2\u6570 –policy：<\/p>\n
uv run scripts\/serve_policy.py –env ALOHA –default_prompt='fold the towel' <\/p>\n
# \u4e3b\u903b\u8f91\u51fd\u6570
\ndef main(args: Args) -> None:
\n policy = create_policy(args) # \u52a0\u8f7d\u7b56\u7565\u5bf9\u8c61
\n policy_metadata = policy.metadata # \u83b7\u53d6\u7b56\u7565\u7684\u5143\u4fe1\u606f<\/p>\n
if args.record: # \u5982\u679c\u542f\u7528\u8bb0\u5f55，\u5305\u88f9\u4e3a\u8bb0\u5f55\u5668
\n policy = _policy.PolicyRecorder(policy, "policy_records")<\/p>\n
# \u83b7\u53d6\u672c\u673a IP \u548c\u4e3b\u673a\u540d（\u65e5\u5fd7\u6253\u5370）
\n hostname = socket.gethostname()
\n local_ip = socket.gethostbyname(hostname)
\n logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)<\/p>\n
# \u521b\u5efa WebSocket \u63a8\u7406\u670d\u52a1\u5bf9\u8c61
\n server = websocket_policy_server.WebsocketPolicyServer(
\n policy=policy,
\n host="0.0.0.0", # \u76d1\u542c\u6240\u6709\u63a5\u53e3
\n port=args.port, # \u6307\u5b9a\u7aef\u53e3
\n metadata=policy_metadata, # \u9644\u52a0\u5143\u4fe1\u606f
\n )
\n server.serve_forever() # \u542f\u52a8\u963b\u585e\u670d\u52a1\u5faa\u73af<\/p>\n
# \u811a\u672c\u5165\u53e3：\u89e3\u6790 CLI \u53c2\u6570\u5e76\u8fd0\u884c\u4e3b\u51fd\u6570
\nif __name__ == "__main__":
\n logging.basicConfig(level=logging.INFO, force=True) # \u521d\u59cb\u5316\u65e5\u5fd7
\n main(tyro.cli(Args)) # \u4f7f\u7528 tyro \u4ece\u547d\u4ee4\u884c\u89e3\u6790 Args \u5e76\u8fd0\u884c <\/p>\n
\u542f\u52a8\u4e00\u4e2a WebSocket \u670d\u52a1\u7aef，\u76d1\u542c\u7aef\u53e3，\u7b49\u5f85\u6765\u81ea\u4e3b\u63a7\u7cfb\u7edf（\u5982 main.py）\u7684\u52a8\u4f5c\u8bf7\u6c42，\u5b8c\u6210\u5982\u4e0b\u529f\u80fd：<\/p>\n
\u2705 \u6700\u7b80\u5355\u7528\u6cd5（\u4f7f\u7528\u9ed8\u8ba4\u6a21\u578b）：<\/p>\n
uv run scripts\/serve_policy.py –env ALOHA –default_prompt='fold the towel' <\/p>\n
\u81ea\u52a8\u4ece S3 \u4e0b\u8f7d pi0_aloha \u7684\u9ed8\u8ba4\u6a21\u578b，\u5e76\u8fd0\u884c\u63a8\u7406\u670d\u52a1\u5668，\u7aef\u53e30.0.0.0:8000<\/p>\n
\u2705 \u4f7f\u7528\u672c\u5730\u6a21\u578b：<\/p>\n
uv run scripts\/serve_policy.py policy:checkpoint \\\\
\n –policy.config=pi0_aloha \\\\
\n –policy.dir=\/home\/yejiangchen\/.cache\/openpi\/openpi-assets\/checkpoints\/pi0_aloha_towel <\/p>\n
\u2705 \u542f\u7528\u63a8\u7406\u884c\u4e3a\u8bb0\u5f55：<\/p>\n
uv run scripts\/serve_policy.py –record … <\/p>\n
\u5c06\u63a8\u7406\u884c\u4e3a\u8bb0\u5f55\u4e0b\u6765，\u4fbf\u4e8e\u8c03\u8bd5<\/p>\n
\u5982\u679c\u5728\u4f7f\u7528\u672c\u5730\u6a21\u578b\u65f6\u51fa\u73b0 bug，\u8fd0\u884c\u4ee5\u4e0b\u547d\u4ee4\u67e5\u770b：<\/p>\n
uv run scripts\/serve_policy.py policy:checkpoint –help <\/p>\n
![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250529123619-683854c3ef61e.png\")
<\/p>\n
\u8fd0\u884c\u6210\u529f\u53ef\u4ee5\u770b\u5230：<\/p>\n
![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250529123620-683854c4803be.png\")
<\/p>\n
\u53ef\u4ee5\u770b\u5230\u65e5\u5fd7：<\/p>\n
\u53ef\u4ee5\u770b\u5230 GPU \u542f\u52a8\u95ee\u9898，\u68c0\u67e5 JAX \u662f\u5426\u7528\u4e0a GPU，\u8f93\u5165：<\/p>\n
(pi0) yejiangchen@yejiangchen-System-Product-Name:~\/Desktop\/Codes\/openpi-main$ python
\nPython 3.11.11 (main, Dec 11 2024, 16:28:39) [GCC 11.2.0] on linux
\nType "help", "copyright", "credits" or "license" for more information.
\n>>> import jax
\n>>> print(jax.devices())
\n[CudaDevice(id=0)] <\/p>\n
\u53ef\u4ee5\u770b\u5230：jax.devices() \u8f93\u51fa\u4e3a [CudaDevice(id=0)]，\u8fd9\u8bf4\u660e JAX \u5728\u5f53\u524d\u73af\u5883\u4e0b\u5df2\u7ecf\u53ef\u4ee5\u6b63\u786e\u8bc6\u522b\u5e76\u8c03\u7528 GPU（CUDA）\u3002\u00a0CUDA \u9a71\u52a8\u3001jax\u3001jaxlib \u5b89\u88c5\u90fd\u6ca1\u6709\u95ee\u9898<\/p>\n
\u5176\u5b9e\u53ea\u8981 jax.devices() \u663e\u793a CudaDevice，\u6a21\u578b\u5b9e\u9645\u8ba1\u7b97\u5c31\u4f1a\u5728 GPU \u4e0a\u6267\u884c，\u5373\u4fbf\u542f\u52a8\u65e5\u5fd7\u6ca1\u9ad8\u4eae\u663e\u793a gpu \u5173\u952e\u5b57\u4e5f\u6ca1\u5173\u7cfb<\/p>\n
JAX\/Flax \u6846\u67b6\u6709\u65f6\u4e0d\u4f1a\u50cf PyTorch \u90a3\u6837\u5728\u65e5\u5fd7\u91cc\u660e\u786e\u8f93\u51fa \u201cUsing CUDA\u201d\u3002\u5b83\u53ea\u4f1a\u5728\u4ee3\u7801\u7b2c\u4e00\u6b21\u6d89\u53ca\u5230\u6570\u7ec4\/\u5f20\u91cf\u8fd0\u7b97\u65f6，\u5c06\u6570\u636e\u79fb\u52a8\u5230 GPU\u3002\u5982\u679c\u5df2\u7ecf\u80fd\u770b\u5230 [CudaDevice(id=0)]，\u8bf4\u660e\u540e\u7eed\u6240\u6709\u8ba1\u7b97\u90fd\u4f1a\u5c3d\u91cf\u5728 GPU \u4e0a\u8fdb\u884c<\/p>\n
\u8fd0\u884c\u65f6\u4f7f\u7528 nvtop \u67e5\u770b GPU \u5360\u7528：<\/p>\n
\u5b89\u88c5：<\/p>\n
sudo apt install nvtop <\/p>\n
\u4f7f\u7528：<\/p>\n
nvtop <\/p>\n
![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250529123621-683854c5d486a.png\")
<\/p>\n
\u53ef\u4ee5\u770b\u5230 GPU \u57fa\u672c\u62c9\u6ee1\u4e86<\/p>\n","protected":false},"excerpt":{"rendered":"
\u6587\u7ae0\u6d4f\u89c8\u9605\u8bfb719\u6b21\uff0c\u70b9\u8d5e24\u6b21\uff0c\u6536\u85cf29\u6b21\u3002\u8be5\u811a\u672c\u5b9e\u73b0\u5c06Pi0\u7b56\u7565\u6a21\u578b\u90e8\u7f72\u4e3aWebSocket\u670d\u52a1\uff0c\u4f9b\u673a\u5668\u4eba\u4e3b\u63a7\u8fdc\u7a0b\u8c03\u7528\u3002\u6838\u5fc3\u6d41\u7a0b\u5305\u62ec\uff1a1) \u901a\u8fc7\u53c2\u6570\u5b9a\u4e49\u73af\u5883\u7c7b\u578b\uff08ALOHA\/DROID\u7b49\uff09\u548c\u6a21\u578b\u8def\u5f84\uff08\u652f\u6301S3\u6216\u672c\u5730\uff09\uff1b2) \u52a0\u8f7d\u9884\u8bad\u7ec3\u7b56\u7565\u6a21\u578b\uff0c\u652f\u6301\u9ed8\u8ba4\u6216\u81ea\u5b9a\u4e49checkpoint\uff1b3) \u542f\u52a8WebSocket\u670d\u52a1\u5668\uff08\u9ed8\u8ba40.0.0.0:8000\uff09\uff0c\u63a5\u6536\u89c2\u6d4b\u6570\u636e\u5e76\u8fd4\u56de\u63a8\u7406\u52a8\u4f5c\u3002<\/p>\n","protected":false},"author":2,"featured_media":40612,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[81,3716,50,384],"topic":[],"class_list":["post-40616","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-server","tag-python","tag-3716","tag-50","tag-384"],"yoast_head":"\n