masato-ka's diary

日々思ったこととか、やったことの備忘録。

【SO-ARM100/101のアプリ実装】5万円で作る！AIロボットアーム SO-ARM100 完全ガイド#5-アプリケーション実装編-

AI AIロボット

1. はじめに

ここまでロボットの組み立てから学習データの収集と学習、学習結果の実行について解説しました。 学習結果の実行ではLeRobotのcontrol_robotスクリプトを使いましたが。この方法では学習したモデルを使いロボットを動かすことができますが、独自のアプリケーションに組み込んで利用することができません。

そこで今回はLeRobotをライブラリとして利用することでオリジナルのアプリケーションを実装する方法について紹介します。

2. LeRobotのアプリケーション実装の全体像

学習させたPolicyを用いてロボットを動かすためには大きく次の処理が必要です。

  1. ロボットへの接続

ロボットのコンフィグファイルを元にロボットのインスタンスを取得し接続します。

  1. ロボットの内部状態、外部センサの状態取得

カメラの画像やロボットの現在の関節値を取得します。

  1. Policyによる推論

2で取得した観測情報を元に学習済みPolicyを使ってロボットの指令値を決定します。

  1. ロボットへの指令値の送信

3で決定した指令値をロボットへ送り、実際にロボットを動かします。

それぞれの処理の実装方法について説明し、最後に全体をまとめたコードを紹介します。

3. ロボットとの接続

LeRobotからロボットへ接続する場合はRobotConfigとManipulatorRobotオブジェクトを使用します。SO-ARM100を接続する場合はSo100RobotConfigを使い、設定を読み込みManipulatoroRobotオブジェクトを作成します。以下のコードはロボットオブジェクトを作成してロボットに接続するコードです。この状態でロボットに対して指示を送ることがあります。

from lerobot.common.robot_devices.robots.configs import So100RobotConfig
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot

config = So100RobotConfig()
config.calibration_dir= <スクリプト実行パスからみたキャリブレーションファイルのパス>

robot = ManipulatorRobot(config)
robot.connect()
## ロボットが急な動きや無理な姿勢、障害物への衝突が発生する可能性があるので
## まだロボットに対して指令値を送らない。

SO100RobotConfigでロボットのコンフィグを読み出します。このconfigは変更可能です。例えば、通常、leaderとfollower両方のアームが読み込まれ接続が行われます。しかし以下のようにSO100RobotConfignのコンストラクタ引数に値を設定することで、leaderアームの設定を無効にし、followerのみ接続するよう変更できます。またcalibration_dirのパスも設定できます。

config = SO100RobotConfig(leader_arms={}, calibration_dir='lerobot/.cache/calibration/so100')
robot = ManipulatorRobot(config)

4. 観測情報の取得

ロボットの関節情報やカメラの画像はManipulatorRobot.capture_observationメソッドからdict型の戻り値として取得できます。カメラはconfigsの設定で指定した値がキー値として設定されます。複数カメラを接続している場合はfor文を使い取り出すことができます。下の例ではリスト内包表記を使用しています。

observation = robot.capture_observation()

state = observation['observation.state']
images = []

image_keys = [key for key in observation if "image" in key]　# カメラのキー値はConfigファイルに設定した値

#画像はnumpy.arrayで格納されている。
for key in image_keys:
    cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
cv2.waitKey(1)

5. ロボットへの指令値の送信

ロボットに対するアクションはrobot.send_actionメソッドを使います。このメソッドは要素0が1軸目の台座の回転に対応している要素数6のリスト(numpy.nbarray or tensor)を引数にとり、そのまま各関節の指令値として与えます。そのため不用意に指令値を与えるとロボットの思わぬ動作を引き起こすことになります。順運動や逆運動学を解き、指令値を決める。もしくは現在の関節角をベースに動作を指定するなど対策をしましょう。

今回はダミーデータとして現在の観測値を取得し、それをそのまま指令値として与えてみます。

if robot.is_connected:
  observation = robot.capture_observation()
  action = observation['observation.state']
  print(action)
  ## 指令値を変更する場合は無理な姿勢にならないか、急な動作を引き起こさないかを確認する。
  ## 警告:適当な値を指令値として与えないこと。
  robot.send_action(action)

5. モデルの読み込みと推論

学習ずみモデルを用いて推論を行う場合はモデルクラスのfrom_pretrained()メソッドに学習したモデルファイルを指定して読み込ませます。ACTを使う場合は次のようにします。

推論結果の戻り値はバッチに従って2次元配列になっています。1次元目の要素は1なので、次元を一つ落として利用します。

from lerobot.common.policies.act.modeling_act import ACTPolicy

policy = ACTPolicy.from_pretrained('モデルファイルまでのパス')
action = policy.select_action(observation) # robot.capture_observationの戻り値
action = action.squeeze(0)#2次元配列から1次元配列へ変換
action = action.to("cpu")
if robot.is_connected:
  robot.send_action(action)

6. 全体のコード

実際にFollowerアームと学習済モデルを使ってロボットを動かす場合は次のようにコードを記述します。以下参考に、ロボットが動き出すトリガーの実装やWEBアプリケーションとの連携などにより、デモの見栄えを良くすることができます。

"""
This is sample code for running the robot controller.
"""


import time

import cv2
import torch

from lerobot.common.policies.act.modeling_act import ACTPolicy
from lerobot.common.robot_devices.robots.configs import So100RobotConfig
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
from lerobot.common.robot_devices.utils import busy_wait

inference_time_s = 60
fps = 30
device = "mps"  # TODO: On Mac, use "mps" or "cpu"

ckpt_path = "masato-ka/act_so100_cls_block_color"
policy = ACTPolicy.from_pretrained(ckpt_path)
policy.to(device)

config = So100RobotConfig(leader_arms={})
config.calibration_dir='lerobot/.cache/calibration/so100'

robot = ManipulatorRobot(config)

if __name__ == '__main__':

    robot.connect()

    try:
        for _ in range(inference_time_s * fps):
            start_time = time.perf_counter()

            # Read the follower state and access the frames from the cameras
            observation = robot.capture_observation()

            image_keys = [key for key in observation if "image" in key]
            for key in image_keys:
                cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
            cv2.waitKey(1)
            # Convert to pytorch format: channel first and float32 in [0,1]
            # with batch dimension
            for name in observation:
                if "image" in name:
                    observation[name] = observation[name].type(torch.float32) / 255
                    observation[name] = observation[name].permute(2, 0, 1).contiguous()
                observation[name] = observation[name].unsqueeze(0)
                observation[name] = observation[name].to(device)

            # Compute the next action with the policy
            # based on the current observation
            action = policy.select_action(observation)
            # Remove batch dimension
            action = action.squeeze(0)
            # Move to cpu, if not already the case
            action = action.to("cpu")
            # Order the robot to move
            robot.send_action(action)

            dt_s = time.perf_counter() - start_time
            busy_wait(1 / fps - dt_s)
    finally:
        robot.disconnect()

7. まとめ

上記コードをベースに、外部イベントによる動作の開始やGUIでの制御といったアプリケーションの構築ができます。LeRobotの付属スクリプの実行だけでなく、WEBブラウザからの入力や他のシステムと連携した動作などインタラクティブなデモアプリケーションを組み上げてください。

次回はより、LeRobotを使いこなすため、Action Chunking Transformerにコンディション(条件)を加えた学習にチャレンジします。例えば、画像上で指定したオブジェクトを取らせたり、言語による作業指示に従って動作させることができるようになります。

▫️画像座標によるコンディショニングの例

x.com

▫️言語によるコンディショニング(VLA)の例

x.com