LLM にアクションを選択させて 犬型ロボット Unitree Go2 を動かしてみた

@dataclass(frozen=True)
class ActionDefinition:
    """Action Registryに登録するアクション定義

    LLMが選択できるアクションを事前定義する。
    velocity / duration は固定値で、APIレスポンスからは変更できない。
    """
    name: str
    description: str
    velocity: float  # m/s（符号で方向を表す: 正=前進, 負=後退）
    duration: float  # 秒

MOVE_FORWARD = ActionDefinition(
    name="move_forward",
    description="1秒間前進する",
    velocity=0.5,
    duration=1.0,
)

MOVE_BACKWARD = ActionDefinition(
    name="move_backward",
    description="1秒間後退する",
    velocity=-0.5,
    duration=1.0,
)

ACTION_REGISTRY: dict[str, ActionDefinition] = {
    MOVE_FORWARD.name: MOVE_FORWARD,
    MOVE_BACKWARD.name: MOVE_BACKWARD,
}

@dataclass(frozen=True)
class ToolInputSchema:
    """Bedrock Converse API の inputSchema"""
    type: str = "object"
    properties: dict[str, Any] = field(default_factory=dict)
    required: list[str] = field(default_factory=list)

@dataclass(frozen=True)
class ToolDefinition:
    """Bedrock Converse API に送信する tool 定義

    ActionDefinition から生成し、API リクエストの tools 引数に渡す。
    """
    name: str
    description: str
    input_schema: ToolInputSchema = field(default_factory=ToolInputSchema)

@dataclass
class ToolUseResponse:
    """Bedrock Converse API の toolUse レスポンスから抽出した情報"""
    tool_use_id: str
    action_name: str

SYSTEM_PROMPT = (
    "あなたは犬型ロボット Unitree Go2 の行動を制御するAIです。"
    "与えられた状況に応じて、最も適切なアクションをひとつ選択してください。"
    "必ずいずれかのアクションを実行してください。"
)

MODEL_ID = "jp.anthropic.claude-sonnet-4-5-20250929-v1:0"

def to_tool_definition(action: ActionDefinition) -> ToolDefinition:
    """ActionDefinition を Bedrock Converse API の ToolDefinition に変換する"""
    return ToolDefinition(
        name=action.name,
        description=action.description,
        input_schema=ToolInputSchema(),
    )

def build_tools_param(
    actions: dict[str, ActionDefinition],
) -> list[dict]:
    """ACTION_REGISTRY から Bedrock Converse API の toolSpec リストを生成する"""
    tools = []
    for action in actions.values():
        td = to_tool_definition(action)
        tools.append(
            {
                "toolSpec": {
                    "name": td.name,
                    "description": td.description,
                    "inputSchema": {
                        "json": {
                            "type": td.input_schema.type,
                            "properties": td.input_schema.properties,
                            "required": td.input_schema.required,
                        }
                    },
                }
            }
        )
    return tools

def parse_tool_use(response: dict) -> ToolUseResponse:
    """Bedrock Converse レスポンスから toolUse ブロックを抽出する"""
    for block in response["output"]["message"]["content"]:
        if "toolUse" in block:
            tool_use = block["toolUse"]
            return ToolUseResponse(
                tool_use_id=tool_use["toolUseId"],
                action_name=tool_use["name"],
            )
    raise ValueError("レスポンスに toolUse ブロックが含まれていません")

def call_agent(
    prompt: str,
    actions: dict[str, ActionDefinition],
) -> ToolUseResponse:
    """Bedrock Converse API を呼び出してアクションを選択させる"""
    client = boto3.client("bedrock-runtime")

    response = client.converse(
        modelId=MODEL_ID,
        system=[{"text": SYSTEM_PROMPT}],
        messages=[
            {
                "role": "user",
                "content": [{"text": prompt}],
            }
        ],
        toolConfig={
            "tools": build_tools_param(actions),
            "toolChoice": {"any": {}},
        },
        inferenceConfig={"maxTokens": 256},
    )

    return parse_tool_use(response)

class RobotController:
    def execute(self, action: ActionDefinition) -> ActionResult:
        """ActionDefinition に基づいてロボットを動かす"""
        try:
            self._sport_client.Move(action.velocity, 0, 0)
            time.sleep(action.duration)
            self._sport_client.StopMove()

            return ActionResult(
                action_name=action.name,
                status=ActionStatus.SUCCESS,
                message=f"{action.name} を実行しました",
            )
        except Exception as e:
            self._sport_client.StopMove()
            return ActionResult(
                action_name=action.name,
                status=ActionStatus.FAILURE,
                message=f"{action.name} の実行に失敗しました: {e}",
            )

def main():
    print("エージェント動作確認（ロボット接続なし）")
    print(f"登録済みアクション: {list(ACTION_REGISTRY.keys())}")

    while True:
        user_input = input("\n指示を入力してください: ")

        print("エージェントに問い合わせ中...")
        response = call_agent(user_input, ACTION_REGISTRY)

        action = ACTION_REGISTRY.get(response.action_name)
        if action is None:
            print(f"未登録のアクション: {response.action_name}")
            continue

        print(f"選択されたアクション: {action.name}")
        print(f"  説明: {action.description}")
        print(f"  速度: {action.velocity} m/s")
        print(f"  時間: {action.duration} 秒")

指示を入力してください: 前に進んで
エージェントに問い合わせ中...
選択されたアクション: move_forward
  説明: 1秒間前進する
  速度: 0.5 m/s
  時間: 1.0 秒

指示を入力してください: 下がって
エージェントに問い合わせ中...
選択されたアクション: move_backward
  説明: 1秒間後退する
  速度: -0.5 m/s
  時間: 1.0 秒

def main():
    network_interface = sys.argv[1] if len(sys.argv) > 1 else None
    controller = RobotController(network_interface)

    print("ロボット制御を開始します。終了するには Ctrl+C を押してください。")

    while True:
        user_input = input("\n指示を入力してください: ")

        print(f"エージェントに問い合わせ中...")
        response = call_agent(user_input, ACTION_REGISTRY)
        print(f"選択されたアクション: {response.action_name}")

        action = ACTION_REGISTRY.get(response.action_name)
        if action is None:
            print(f"未登録のアクション: {response.action_name}")
            continue

        result = controller.execute(action)
        print(f"結果: {result.status.value} - {result.message}")