Go to Qiita Advent Calendar 2024 Top
LoginSignup
3
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@hbvcg00(Tatsuki Nishimura)

LeRobotを使ってGymnasiumで作成した自作環境でDiffusion Policyを回すまでのメモ

Last updated at Posted at 2024-11-01

最初に

最近LeRobotが話題ですが，まだ新しいゆえに日本語での記事があまりなかったので，LeRobotの導入→自作環境に適用→データセット作成→学習を回すまでの流れをメモ程度にまとめてみました．かなり自己流でやってるので本家のやり方とは異なっているかもです．間違っているところあれば教えてください．

注意
本記事ではstate-basedのDiffusion Policyを取り扱っています．image-basedの方は対応してないです．

LeRobotとは

HuggingFaceが管理している，ロボットの模倣学習や強化学習への応用をしやすくするフレームワーク（という認識）で，HuggingFaceからモデルやデータセットを簡単にダウンロードして実行できる機能や，事前にDiffusion PolicyやACTなどの有名な模倣学習アルゴリズムが実装されていてすぐに使えるという点が便利です．

2024年11月現在は活発に更新されており，exampleも段々と充実してきています．ですが，実際自分が研究で使う環境に適用するまでのチュートリアルというのはあまり多くない印象です．ましてや日本語での記事はもっと見かけないです．自分は現在研究でLeRobotを使っていて結構苦労したので，メモ程度に情報を残しておきます．

インストール

基本はREADMEのインストールに従えば問題ないです．
https://github.com/huggingface/lerobot/blob/main/README.md#installation

私の場合はcondaで環境構築するのがあまり好みではなかったので，Poetryで環境構築しました．LeRobotのルートディレクトリでpoetry installすればそれだけで済みます．

それと，wandbのアカウントは事前に作っておいたほうが便利です．Web上で学習状況を確認できるようになります．

Gymnasium環境の追加

自分で作成したGymnasiumの環境をLeRobotで認識させます．今回は例としてtestという名前の環境を追加してみます．

まず，testのディレクトリ構成はこんな感じです．gym-pushtを名前変えただけです．

test
gym_test
├── __init__.py               # 環境を登録するregister関数が入ってる
└── envs                       
    ├── __init__.py            
    ├── hoge.py            
    └── fuga.py

注意
今の所環境名はgym_〇〇にしといたほうが無難っぽいです．↓参照
https://github.com/huggingface/lerobot/blob/e0df56de621b6f7ee501719ee0b1e4af00a98635/lerobot/common/envs/factory.py#L33

このgym_testを，LeRobotの一番上のディレクトリにとりあえず配置します．（本当は良くない気がしますが）

lerobot/
├── gym_test/ # とりあえずここに配置
├── examples/                 
├── lerobot/                    
├── docker/            
├── tests/
├── media/
├── pyproject.toml
.

基本的にgym_test直下の__init__.pyが最初に読み込まれるようにすればOKなので，lerobot/__init__.pyを編集します．

1. 環境をimport

lerobot/__init__.py(一部抜粋)
import gym_test

2. available_tasks_per_envを更新

lerobot/__init__.py(一部抜粋)
available_tasks_per_env = {
    "aloha": [
        "AlohaInsertion-v0",
        "AlohaTransferCube-v0",
    ],
    "pusht": ["PushT-v0"],
    "xarm": ["XarmLift-v0"],
    "dora_aloha_real": ["DoraAloha-v0", "DoraKoch-v0", "DoraReachy2-v0"],
    "test": ["Test-v0"],   # 今回追加
}

ここで，keyには環境名，valueにはタスク名を入れます．
こうすることで，gym_(環境名)/(タスク名)の形にして環境を引っ張ってこれるんですね．
（参照）
https://github.com/huggingface/lerobot/blob/e0df56de621b6f7ee501719ee0b1e4af00a98635/lerobot/common/envs/factory.py#L43

gym_test/__init__.py
register(
    id="gym_test/Test-v0",
    entry_point="...",
    max_episode_steps=...,
    reward_threshold=...,
)

3. available_datasets_per_envも更新

lerobot/__init__.py(一部抜粋)
available_datasets_per_env = {
    "aloha": [...],
    "pusht": [...],
    "xarm": [...],
    "dora_aloha_real": [...],
    "test": ["custom/test"], # 今回追加
}

ここではkeyに環境名，valueには今後使うrepo_idを登録します．repo_idは任意に決められるっぽいです．

以上で環境の追加は終わりです．

データセットの作成

次にデータセットを作成します．LeRobotではHuggingFace上から落としてきたデータセットを取り込む方法は多く解説されていますが，手元でデータセットを作る方法はあんまり詳しくない印象です．

まず，LeRobotで使用されるデータセットはLeRobotDatasetというクラスであり，これについては
https://github.com/huggingface/lerobot?tab=readme-ov-file#the-lerobotdataset-format
で公式に詳しく解説されていますので，こちらを読むことをおすすめします．
要はPytorchのDatasetをLeRobot側でラップしてます．
LeRobotDatasetの最小構成は以下のようになります．

LeRobotDataset
LeRobotDataset
  ├ hf_dataset
  ├ episode_data_index
  ├ stats
  ├ info

とりあえずこの4つを揃えてLeRobotDataset形式に持っていくことができればOKそうです．

作り方はhttps://github.com/huggingface/lerobot/blob/main/lerobot/common/datasets/push_dataset_to_hub/pusht_zarr_format.py
あたりを大いに参考にしています．

データセットの作るだいたいのコードを載せておきます．私の環境でしか動かないので，いい感じにimportとかつけて自分の環境用に編集してください．

データセットを作る関数まとめ

import pickle
from lerobot.common.datasets.compute_stats import compute_stats
import torch
from datasets import Dataset, Features, Image, Sequence, Value
from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
from lerobot.common.datasets.push_dataset_to_hub.utils import concatenate_episodes
from lerobot.common.datasets.utils import (
    calculate_episode_data_index,
    hf_transform_to_torch,
)
from pathlib import Path
from lerobot.scripts.push_dataset_to_hub import save_meta_data

def load_from_raw(path: Path, fps: int):
    data = None
    # 事前に用意したpickle形式のエキスパートデータを読み込む
    with open(path, "rb") as f:
        data = pickle.load(f)
    ep_dicts = []
    for episode_index, ts in enumerate(data):
        ep_dict = {
            "observation.state": None,
            "observation.environment_state": None,
            "action": None,
            "episode_index": None,
            "frame_index": None,
            "timestamp": None,
            "next.reward": None,
            "next.done": None,
            "next.success": None,
        }
        num_frames = len(ts.obs)
        ep_dict["observation.state"] = torch.tensor(ts.obs)
        # observation.environment_stateに入れる値はこれであってるかわからない
        ep_dict["observation.environment_state"] = torch.tensor(ts.obs)
        ep_dict["action"] = torch.tensor(ts.acts)
        ep_dict["episode_index"] = torch.tensor(
            [episode_index] * num_frames, dtype=torch.int64
        )
        ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
        ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
        ep_dict["next.reward"] = torch.zeros(num_frames, dtype=torch.float32)
        ep_dict["next.done"] = torch.zeros(num_frames, dtype=torch.bool)
        ep_dict["next.done"][-1] = True
        ep_dict["next.success"] = torch.zeros(num_frames, dtype=torch.bool)
        ep_dicts.append(ep_dict)
    data_dict = concatenate_episodes(ep_dicts)
    total_frames = data_dict["frame_index"].shape[0]
    data_dict["index"] = torch.arange(0, total_frames, 1)
    return data_dict


def to_hf_dataset(data_dict):
    features = {}
    features["observation.state"] = Sequence(
        length=data_dict["observation.state"].shape[1],
        feature=Value(dtype="float32", id=None),
    )
    features["observation.environment_state"] = Sequence(
        length=data_dict["observation.environment_state"].shape[1],
        feature=Value(dtype="float32", id=None),
    )
    features["action"] = Sequence(
        length=data_dict["action"].shape[1], feature=Value(dtype="int64", id=None)
    )
    features["episode_index"] = Value(dtype="int64", id=None)
    features["frame_index"] = Value(dtype="int64", id=None)
    features["timestamp"] = Value(dtype="float32", id=None)
    features["next.reward"] = Value(dtype="float32", id=None)
    features["next.done"] = Value(dtype="bool", id=None)
    features["index"] = Value(dtype="int64", id=None)
    features["next.success"] = Value(dtype="bool", id=None)
    features["action_is_pad"] = Value(dtype="bool", id=None)

    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
    hf_dataset.set_transform(hf_transform_to_torch)
    return hf_dataset


def from_raw_to_lerobot_format(
    path: Path,
    fps: int,
):
    data_dict = load_from_raw(path=path, fps=fps)
    hf_dataset = to_hf_dataset(data_dict)
    episode_data_index = calculate_episode_data_index(hf_dataset)
    info = {
        "codebase_version": CODEBASE_VERSION,
        "fps": fps,
    }
    return hf_dataset, episode_data_index, info


def get_dataset(
    repo_id: str, root: str, raw_data_path: str, fps: int, save=False
) -> LeRobotDataset:
    local_dir = Path(root) / repo_id

    if save:
        import shutil

        if local_dir.exists():
            shutil.rmtree(local_dir)
        episodes_dir = local_dir / "episodes"
        episodes_dir.mkdir(parents=True, exist_ok=True)

    hf_dataset, episode_data_index, info = from_raw_to_lerobot_format(
        path=raw_data_path, fps=fps
    )
    lerobot_dataset = LeRobotDataset.from_preloaded(
        repo_id=repo_id,
        hf_dataset=hf_dataset,
        episode_data_index=episode_data_index,
        info=info,
    )
    stats = compute_stats(lerobot_dataset)
    lerobot_dataset.stats = stats
    if save:
        hf_dataset = hf_dataset.with_format(
            None
        )  # to remove transforms that cant be saved
        hf_dataset.save_to_disk(str(local_dir / "train"))
        meta_data_dir = local_dir / "meta_data"
        save_meta_data(info, stats, episode_data_index, meta_data_dir)
    return lerobot_dataset


if __name__ == "__main__":
    repo_id = "custom/test"
    root = "custom_datasets"
    raw_data_path = (pickle形式のexpert_dataのpath)
    fps = ...
    dataset = get_dataset(
        repo_id=repo_id, root=root, raw_data_path=raw_data_path, fps=fps, save=True
    )

それぞれの関数について簡単に解説

  • load_from_raw
    • pickle形式のエキスパートデータを整形
    • ep_dictには各エピソードごとの情報が入る
    • 自分が持ってるエキスパートデータの形式に合わせてep_dictに入れる
    • data_dictを出力
  • to_hf_dataset
    • data_dictをHuggingFace Dataset形式に変換
    • HuggingFace DatasetはHuggingFaceからダウンロードできるデータセットの形式のこと？
    • この関数でhf_datasetが得られる
  • from_raw_to_lerobot_format
    • hf_datasetの情報を基にepisode_data_indexとinfoを追加
  • get_dataset
    • LeRobotDatasetを得る関数
    • 作成したLeRobotDatasetはローカルに保存
    • statsもここで計算

という感じです．
mainの部分を実行すればローカルでデータセットが保存されます．

注意
こういう構成にしたらとりあえずデータセットできた！みたいなレベルのものなので公式の想定している作り方ではないかもしれないです．

学習を回す

作成したデータセットを用いてDiffusion Policyアルゴリズムで学習してみます．
公式チュートリアルの3_train_policy.pyや4_train_policy_with_script.mdで解説されている内容を参考にしています．

今回はチュートリアル4の内容であるコマンドラインからlerobot/train.pyを叩く方法に従います．

パラメータ設定

gym-test環境のパラメータ設定と，Diffusion Policyのパラメータ設定をそれぞれ行います.
まずはlerobot/configs/env/test.yamlを追加し，以下を書き込みます．

lerobot/configs/env/test.yaml
# @package _global_

fps: 100

env:
  name: test
  task: Test-v0
  state_dim: 2
  action_dim: 2
  fps: ${fps}
  episode_length: 300

state_dimとaction_Dim，fpsは自分の環境に合わせてください．

次に，まずはlerobot/configs/policy/diffusion_test.yamlを追加し，以下を書き込みます．

lerobot/configs/policy/diffusion_test.yaml
hydra:
  run:
    # Set `dir` to where you would like to save all of the run outputs. If you run another training session
    # with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
    dir: outputs/train/${now:%Y-%m-%d}/${now:%H-%M-%S}_${env.name}_${policy.name}_${hydra.job.name}
  job:
    name: test_${now:%Y-%m-%d}_${now:%H-%M-%S}

resume: false
device: cuda  # cpu
use_amp: true
seed: 100000
dataset_repo_id: custom/test

training:
  offline_steps: 200000
  online_steps: 0
  eval_freq: 10000
  save_freq: 10000
  log_freq: 250
  save_checkpoint: true
  num_workers: 4
  batch_size: 256
  grad_clip_norm: 10
  lr: 1.0e-4
  lr_scheduler: cosine
  lr_warmup_steps: 500
  adam_betas: [0.95, 0.999]
  adam_eps: 1.0e-8
  adam_weight_decay: 1.0e-6
  online_steps_between_rollouts: 1
  delta_timestamps:
    observation.environment_state: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]"
    observation.state: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]"
    action: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1 - ${policy.n_obs_steps} + ${policy.horizon})]"
  # The original implementation doesn't sample frames for the last 7 steps,
  # which avoids excessive padding and leads to improved training results.
  drop_n_last_frames: 7  # ${policy.horizon} - ${policy.n_action_steps} - ${policy.n_obs_steps} + 1

eval:
  n_episodes: 50
  batch_size: 4

policy:
  name: diffusion
  n_obs_steps: 2
  horizon: 16
  n_action_steps: 8

  input_shapes:
    observation.environment_state: [2]
    observation.state: ["${env.state_dim}"]
  output_shapes:
    action: ["${env.action_dim}"]

  # Normalization / Unnormalization
  input_normalization_modes:
    observation.environment_state: mean_std
    observation.state: min_max
  output_normalization_modes:
    action: min_max

  # Architecture / modeling.
  # Vision backbone.
  vision_backbone: resnet18
  crop_shape: [84, 84]
  crop_is_random: True
  pretrained_backbone_weights: null
  use_group_norm: True
  spatial_softmax_num_keypoints: 32
  # Unet.
  down_dims: [512, 1024, 2048]
  kernel_size: 5
  n_groups: 8
  diffusion_step_embed_dim: 128
  use_film_scale_modulation: True
  # Noise scheduler.
  noise_scheduler_type: DDIM
  num_train_timesteps: 100
  beta_schedule: squaredcos_cap_v2
  beta_start: 0.0001
  beta_end: 0.02
  prediction_type: epsilon # epsilon / sample
  clip_sample: True
  clip_sample_range: 1.0

  # Inference
  # num_inference_steps: 10  # if not provided, defaults to `num_train_timesteps`

  # Loss computation
  do_mask_loss_for_padding: false

wandb:
  enable: true
  # Set to true to disable saving an artifact despite save_checkpoint == True
  disable_artifact: true
  project: test
  notes: ""

trainingやpolicyのパラメータは自分好みに調整してください．
wandbのdisable_artifactはtrueにしておかないと学習の途中でモデルをwandb上にバックアップしてしまい，無料ユーザの容量制限をすぐに超えてしまいました．

学習

DATA_DIR="(データセットを保存した場所)" python lerobot/scripts/train.py policy=diffusion_test env=test

で学習を実行できます．結果はdiffusion_test.yamlのhydra.run.dirで指定したディレクトリに生成されます．
また，wandbにログインしておけば学習の進捗も確認できます．

また，便利機能としてresumeというものがあります．
https://github.com/huggingface/lerobot/blob/main/examples/5_resume_training.md
でも解説されている通り，学習が途中でストップしてもそこから開始できます．
dataloaderが落ちたりしても再開できるのでとても便利です．

評価

https://github.com/huggingface/lerobot/blob/main/examples/2_evaluate_pretrained_policy.py
を少し編集すればできます．

2_evaluate_pretrained_policy.py(一部抜粋)
# Download the diffusion policy for pusht environment
# pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))
# OR uncomment the following to evaluate a policy from the local outputs/train folder.
 pretrained_policy_path = Path("outputs/train/2024-10-19/12-34-50_test_diffusion_test_2024-10-19_12-34-50/checkpoints/150000/pretrained_model")
2_evaluate_pretrained_policy.py(一部抜粋)
env = gym.make(
    "gym_test/Test-v0",
    max_episode_steps=300,
)

みたいな感じで書き換えればあとは評価できます．

まとめ

LeRobotのインストール→自作環境の導入→データセットの作成→学習→評価
までの流れを書いてみました．誰かの役にたてることを願っています．
間違い等あったら教えてください．

3
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
3
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?