Genesis 入門 (11) - 強化学習による移動ポリシーの学習 の コード解説

強化学習による移動ポリシーの学習のコード解説をまとめました。

前回

1. 強化学習による移動ポリシーの学習のコード

以下で使用したコードの解説をまとめました。

次の3つのコードで構成されています。

2. go2_env.py

四足歩行ロボット「Uitree Go2」の強化学習環境を実装したコードです。

・環境の構成

・観測空間

・行動空間

・報酬

・終了条件

コードは、次のとおりです。

import torch
import math
import genesis as gs
from genesis.utils.geom import quat_to_xyz, transform_by_quat, inv_quat, transform_quat_by_quat


# 指定範囲でランダムな浮動小数を取得
def gs_rand_float(lower, upper, shape, device):
    return (upper - lower) * torch.rand(size=shape, device=device) + lower

# 環境
class Go2Env:
    # 初期化
    def __init__(self, num_envs, env_cfg, obs_cfg, reward_cfg, command_cfg, show_viewer=False, device="cuda"):
        # デバイス設定
        self.device = torch.device(device)

        # 環境パラメータの設定
        self.num_envs = num_envs  # 環境数
        self.num_obs = obs_cfg["num_obs"]  # 観測値の数
        self.num_privileged_obs = None  # 特権観測値 (未使用)
        self.num_actions = env_cfg["num_actions"]  # 行動の数
        self.num_commands = command_cfg["num_commands"]  # 命令の数

        self.simulate_action_latency = True  # 実ロボットでの1ステップ遅延をシミュレート
        self.dt = 0.02  # 制御周期 (50Hz)
        self.max_episode_length = math.ceil(env_cfg["episode_length_s"] / self.dt)  # エピソード長

        # 各種設定の保存
        self.env_cfg = env_cfg
        self.obs_cfg = obs_cfg
        self.reward_cfg = reward_cfg
        self.command_cfg = command_cfg

        self.obs_scales = obs_cfg["obs_scales"]  # 観測値スケール
        self.reward_scales = reward_cfg["reward_scales"]  # 報酬スケール

        # シーンの作成
        self.scene = gs.Scene(
            sim_options=gs.options.SimOptions(dt=self.dt, substeps=2),
            viewer_options=gs.options.ViewerOptions(
                max_FPS=int(0.5 / self.dt),
                camera_pos=(2.0, 0.0, 2.5),
                camera_lookat=(0.0, 0.0, 0.5),
                camera_fov=40,
            ),
            vis_options=gs.options.VisOptions(n_rendered_envs=1),
            rigid_options=gs.options.RigidOptions(
                dt=self.dt,
                constraint_solver=gs.constraint_solver.Newton,
                enable_collision=True,
                enable_joint_limit=True,
            ),
            show_viewer=show_viewer,
        )

        # 地面の追加
        self.scene.add_entity(gs.morphs.URDF(file="urdf/plane/plane.urdf", fixed=True))

        # ロボットの追加
        self.base_init_pos = torch.tensor(self.env_cfg["base_init_pos"], device=self.device)  # 初期位置
        self.base_init_quat = torch.tensor(self.env_cfg["base_init_quat"], device=self.device)  # 初期回転
        self.inv_base_init_quat = inv_quat(self.base_init_quat)  # 初期回転の逆クォータニオン
        self.robot = self.scene.add_entity(
            gs.morphs.URDF(
                file="urdf/go2/urdf/go2.urdf",
                pos=self.base_init_pos.cpu().numpy(),
                quat=self.base_init_quat.cpu().numpy(),
            ),
        )

        # シーンのビルド
        self.scene.build(n_envs=num_envs)

        # 関節名からインデックスを取得
        self.motor_dofs = [self.robot.get_joint(name).dof_idx_local for name in self.env_cfg["dof_names"]]

        # PDコントローラのパラメータの設定
        self.robot.set_dofs_kp([self.env_cfg["kp"]] * self.num_actions, self.motor_dofs)
        self.robot.set_dofs_kv([self.env_cfg["kd"]] * self.num_actions, self.motor_dofs)

        # 報酬関数の準備と報酬スケールの時間ステップでの調整
        self.reward_functions, self.episode_sums = dict(), dict()
        for name in self.reward_scales.keys():
            self.reward_scales[name] *= self.dt
            self.reward_functions[name] = getattr(self, "_reward_" + name)
            self.episode_sums[name] = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_float)

        # バッファの初期化
        self.base_lin_vel = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)  # 基本線速度
        self.base_ang_vel = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)  # 基本角速度
        self.projected_gravity = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)  # 投影重力
        self.global_gravity = torch.tensor([0.0, 0.0, -1.0], device=self.device, dtype=gs.tc_float).repeat(
            self.num_envs, 1
        )  # 全体の重力
        self.obs_buf = torch.zeros((self.num_envs, self.num_obs), device=self.device, dtype=gs.tc_float)  # 観測バッファ
        self.rew_buf = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_float)  # 報酬バッファ
        self.reset_buf = torch.ones((self.num_envs,), device=self.device, dtype=gs.tc_int)  # リセットバッファ
        self.episode_length_buf = torch.zeros((self.num_envs,), device=self.device, dtype=gs.tc_int)  # エピソード長バッファ
        self.commands = torch.zeros((self.num_envs, self.num_commands), device=self.device, dtype=gs.tc_float)  # 命令
        self.commands_scale = torch.tensor(
            [self.obs_scales["lin_vel"], self.obs_scales["lin_vel"], self.obs_scales["ang_vel"]],
            device=self.device,
            dtype=gs.tc_float,
        )  # 命令スケール
        self.actions = torch.zeros((self.num_envs, self.num_actions), device=self.device, dtype=gs.tc_float)  # 行動
        self.last_actions = torch.zeros_like(self.actions)  # 前回の行動
        self.dof_pos = torch.zeros_like(self.actions)  # 関節位置
        self.dof_vel = torch.zeros_like(self.actions)  # 関節速度
        self.last_dof_vel = torch.zeros_like(self.actions)  # 前回の関節速度
        self.base_pos = torch.zeros((self.num_envs, 3), device=self.device, dtype=gs.tc_float)  # 基本位置
        self.base_quat = torch.zeros((self.num_envs, 4), device=self.device, dtype=gs.tc_float)  # 基本回転
        self.default_dof_pos = torch.tensor(
            [self.env_cfg["default_joint_angles"][name] for name in self.env_cfg["dof_names"]],
            device=self.device,
            dtype=gs.tc_float,
        )  # デフォルト関節位置
        self.extras = dict()  # ロギング用の追加情報

    # 指定された環境インデックスで新しい命令を再サンプリング
    def _resample_commands(self, envs_idx):
        self.commands[envs_idx, 0] = gs_rand_float(*self.command_cfg["lin_vel_x_range"], (len(envs_idx),), self.device)
        self.commands[envs_idx, 1] = gs_rand_float(*self.command_cfg["lin_vel_y_range"], (len(envs_idx),), self.device)
        self.commands[envs_idx, 2] = gs_rand_float(*self.command_cfg["ang_vel_range"], (len(envs_idx),), self.device)

    # ステップ
    def step(self, actions):
        # 行動を実行して環境を更新
        self.actions = torch.clip(actions, -self.env_cfg["clip_actions"], self.env_cfg["clip_actions"])
        exec_actions = self.last_actions if self.simulate_action_latency else self.actions
        target_dof_pos = exec_actions * self.env_cfg["action_scale"] + self.default_dof_pos
        self.robot.control_dofs_position(target_dof_pos, self.motor_dofs)
        self.scene.step()

        # バッファの更新
        self.episode_length_buf += 1
        self.base_pos[:] = self.robot.get_pos()
        self.base_quat[:] = self.robot.get_quat()
        self.base_euler = quat_to_xyz(
            transform_quat_by_quat(torch.ones_like(self.base_quat) * self.inv_base_init_quat, self.base_quat)
        )
        inv_base_quat = inv_quat(self.base_quat)
        self.base_lin_vel[:] = transform_by_quat(self.robot.get_vel(), inv_base_quat)
        self.base_ang_vel[:] = transform_by_quat(self.robot.get_ang(), inv_base_quat)
        self.projected_gravity = transform_by_quat(self.global_gravity, inv_base_quat)
        self.dof_pos[:] = self.robot.get_dofs_position(self.motor_dofs)
        self.dof_vel[:] = self.robot.get_dofs_velocity(self.motor_dofs)

        # 命令を再サンプリング
        envs_idx = (
            (self.episode_length_buf % int(self.env_cfg["resampling_time_s"] / self.dt) == 0)
            .nonzero(as_tuple=False)
            .flatten()
        )
        self._resample_commands(envs_idx)

        # 終了チェック
        self.reset_buf = self.episode_length_buf > self.max_episode_length
        self.reset_buf |= torch.abs(self.base_euler[:, 1]) > self.env_cfg["termination_if_pitch_greater_than"]
        self.reset_buf |= torch.abs(self.base_euler[:, 0]) > self.env_cfg["termination_if_roll_greater_than"]

        # タイムアウトの処理
        time_out_idx = (self.episode_length_buf > self.max_episode_length).nonzero(as_tuple=False).flatten()
        self.extras["time_outs"] = torch.zeros_like(self.reset_buf, device=self.device, dtype=gs.tc_float)
        self.extras["time_outs"][time_out_idx] = 1.0

        # リセットが必要な環境をリセット
        self.reset_idx(self.reset_buf.nonzero(as_tuple=False).flatten())

        # 報酬を計算
        self.rew_buf[:] = 0.0
        for name, reward_func in self.reward_functions.items():
            # 各報酬関数を実行し、スケールを適用
            rew = reward_func() * self.reward_scales[name]
            self.rew_buf += rew
            self.episode_sums[name] += rew

        # 観測値の計算
        self.obs_buf = torch.cat(
            [
                self.base_ang_vel * self.obs_scales["ang_vel"],  # 3
                self.projected_gravity,  # 3
                self.commands * self.commands_scale,  # 3
                (self.dof_pos - self.default_dof_pos) * self.obs_scales["dof_pos"],  # 12
                self.dof_vel * self.obs_scales["dof_vel"],  # 12
                self.actions,  # 12
            ],
            axis=-1,
        )

        # 前回の行動と関節速度を更新
        self.last_actions[:] = self.actions[:]
        self.last_dof_vel[:] = self.dof_vel[:]

        return self.obs_buf, None, self.rew_buf, self.reset_buf, self.extras

    # 観測の取得
    def get_observations(self):
        return self.obs_buf

    # 特権的な観測の取得
    def get_privileged_observations(self):
        return None


    # 指定環境のリセット
    def reset_idx(self, envs_idx):
        if len(envs_idx) == 0:
            return

        # dofのリセット
        self.dof_pos[envs_idx] = self.default_dof_pos  # デフォルトの関節位置に設定
        self.dof_vel[envs_idx] = 0.0  # 関節速度をゼロにリセット
        self.robot.set_dofs_position(
            position=self.dof_pos[envs_idx],
            dofs_idx_local=self.motor_dofs,
            zero_velocity=True,
            envs_idx=envs_idx,
        )

        # baseのリセット
        self.base_pos[envs_idx] = self.base_init_pos
        self.base_quat[envs_idx] = self.base_init_quat.reshape(1, -1)
        self.robot.set_pos(self.base_pos[envs_idx], zero_velocity=False, envs_idx=envs_idx)
        self.robot.set_quat(self.base_quat[envs_idx], zero_velocity=False, envs_idx=envs_idx)
        self.base_lin_vel[envs_idx] = 0
        self.base_ang_vel[envs_idx] = 0
        self.robot.zero_all_dofs_velocity(envs_idx)

        # バッファのリセット
        self.last_actions[envs_idx] = 0.0
        self.last_dof_vel[envs_idx] = 0.0
        self.episode_length_buf[envs_idx] = 0
        self.reset_buf[envs_idx] = True

        # 補助データの初期化
        self.extras["episode"] = {}
        for key in self.episode_sums.keys():
            # 報酬関連のデータを初期化
            self.extras["episode"]["rew_" + key] = (
                torch.mean(self.episode_sums[key][envs_idx]).item() / self.env_cfg["episode_length_s"]
            )
            self.episode_sums[key][envs_idx] = 0.0

        self._resample_commands(envs_idx)

    # リセット
    def reset(self):
        # すべての環境をリセット
        self.reset_buf[:] = True
        self.reset_idx(torch.arange(self.num_envs, device=self.device))
        return self.obs_buf, None

    # 直線速度命令の追従（xy軸）
    def _reward_tracking_lin_vel(self):
        # Tracking of linear velocity commands (xy axes)
        lin_vel_error = torch.sum(torch.square(self.commands[:, :2] - self.base_lin_vel[:, :2]), dim=1)
        return torch.exp(-lin_vel_error / self.reward_cfg["tracking_sigma"])

    # 角速度命令の追従（yaw）
    def _reward_tracking_ang_vel(self):
        ang_vel_error = torch.square(self.commands[:, 2] - self.base_ang_vel[:, 2])
        return torch.exp(-ang_vel_error / self.reward_cfg["tracking_sigma"])

    # z軸方向の線形速度のペナルティ化
    def _reward_lin_vel_z(self):
        return torch.square(self.base_lin_vel[:, 2])

    # 行動の変化率のペナルティ化
    def _reward_action_rate(self):
        return torch.sum(torch.square(self.last_actions - self.actions), dim=1)

    # 姿勢がデフォルトから大きく外れた場合のペナルティ化
    def _reward_similar_to_default(self):
        return torch.sum(torch.abs(self.dof_pos - self.default_dof_pos), dim=1)

    # ベース高さが目標値から外れた場合のペナルティ化
    def _reward_base_height(self):
        return torch.square(self.base_pos[:, 2] - self.reward_cfg["base_height_target"])

3. go2_train.py

強化学習環境 (go2_env.py) を使用して学習するコードです。

コードは、次のとおりです。

import argparse
import os
import pickle
import shutil
from go2_env import Go2Env
from rsl_rl.runners import OnPolicyRunner
import genesis as gs

# 学習設定の取得
def get_train_cfg(exp_name, max_iterations):
    train_cfg_dict = {
        "algorithm": {  # アルゴリズムの設定
            "clip_param": 0.2,  # クリッピングのパラメータ
            "desired_kl": 0.01,  # KLダイバージェンスの目標値
            "entropy_coef": 0.01,  # エントロピーの係数
            "gamma": 0.99,  # 割引率
            "lam": 0.95,  # GAEのパラメータ
            "learning_rate": 0.001,  # 学習率
            "max_grad_norm": 1.0,  # 勾配の最大値
            "num_learning_epochs": 5,  # エポック数
            "num_mini_batches": 4,  # ミニバッチの数
            "schedule": "adaptive",  # スケジュール方式
            "use_clipped_value_loss": True,  # 価値関数損失のクリッピングを使用するか
            "value_loss_coef": 1.0,  # 価値関数の損失係数
        },
        "init_member_classes": {},  # メンバークラスの初期化設定
        "policy": {  # ポリシーの設定
            "activation": "elu",  # 活性化関数
            "actor_hidden_dims": [512, 256, 128],  # Actorネットワークの隠れ層サイズ
            "critic_hidden_dims": [512, 256, 128],  # Criticネットワークの隠れ層サイズ
            "init_noise_std": 1.0,  # 初期ノイズ標準偏差
        },
        "runner": {  # ランナーの設定
            "algorithm_class_name": "PPO",  # 使用するアルゴリズム名
            "checkpoint": -1,  # チェックポイントのインデックス
            "experiment_name": exp_name,  # 実験名
            "load_run": -1,  # 実行のロード設定
            "log_interval": 1,  # ログの間隔
            "max_iterations": max_iterations,  # 最大反復回数
            "num_steps_per_env": 24,  # 環境ごとのステップ数
            "policy_class_name": "ActorCritic",  # ポリシークラス名
            "record_interval": -1,  # 記録間隔
            "resume": False,  # 再開設定
            "resume_path": None,  # 再開パス
            "run_name": "",  # 実行名
            "runner_class_name": "runner_class_name",  # ランナークラス名
            "save_interval": 100,  # 保存間隔
        },
        "runner_class_name": "OnPolicyRunner",  # ランナーのクラス名
        "seed": 1,  # ランダムシード
    }

    return train_cfg_dict


# 設定の取得
def get_cfgs():
    env_cfg = {
        "num_actions": 12,  # 行動の数
        "default_joint_angles": {  # 各関節のデフォルト角度 [ラジアン]
            "FL_hip_joint": 0.0,
            "FR_hip_joint": 0.0,
            "RL_hip_joint": 0.0,
            "RR_hip_joint": 0.0,
            "FL_thigh_joint": 0.8,
            "FR_thigh_joint": 0.8,
            "RL_thigh_joint": 1.0,
            "RR_thigh_joint": 1.0,
            "FL_calf_joint": -1.5,
            "FR_calf_joint": -1.5,
            "RL_calf_joint": -1.5,
            "RR_calf_joint": -1.5,
        },
        "dof_names": [  # 自由度（Degree of Freedom）の名前
            "FR_hip_joint",
            "FR_thigh_joint",
            "FR_calf_joint",
            "FL_hip_joint",
            "FL_thigh_joint",
            "FL_calf_joint",
            "RR_hip_joint",
            "RR_thigh_joint",
            "RR_calf_joint",
            "RL_hip_joint",
            "RL_thigh_joint",
            "RL_calf_joint",
        ],
        "kp": 20.0,  # 比例ゲイン
        "kd": 0.5,  # 微分ゲイン
        "termination_if_roll_greater_than": 10,  # 転倒判定（ロール角度）
        "termination_if_pitch_greater_than": 10,  # 転倒判定（ピッチ角度）
        "base_init_pos": [0.0, 0.0, 0.42],  # ベースの初期位置
        "base_init_quat": [0.0, 0.0, 0.0, 1.0],  # ベースの初期クォータニオン
        "episode_length_s": 20.0,  # エピソードの長さ（秒）
        "resampling_time_s": 4.0,  # サンプリング間隔（秒）
        "action_scale": 0.25,  # アクションのスケール
        "simulate_action_latency": True,  # アクション遅延をシミュレートするか
        "clip_actions": 100.0,  # アクションのクリッピング
    }
    obs_cfg = {
        "num_obs": 45,  # 観測値の数
        "obs_scales": {  # 各観測値のスケール
            "lin_vel": 2.0,  # 線速度
            "ang_vel": 0.25,  # 角速度
            "dof_pos": 1.0,  # 自由度の位置
            "dof_vel": 0.05,  # 自由度の速度
        },
    }
    reward_cfg = {
        "tracking_sigma": 0.25,  # トラッキング誤差の標準偏差
        "base_height_target": 0.3,  # ベースの目標高さ
        "feet_height_target": 0.075,  # 足の目標高さ
        "reward_scales": {  # 報酬のスケール
            "tracking_lin_vel": 1.0,  # 線速度の追跡
            "tracking_ang_vel": 0.2,  # 角速度の追跡
            "lin_vel_z": -1.0,  # Z方向の速度
            "base_height": -50.0,  # ベースの高さ
            "action_rate": -0.005,  # アクション変化率
            "similar_to_default": -0.1,  # デフォルトへの類似性
        },
    }
    command_cfg = {
        "num_commands": 3,  # 命令の数
        "lin_vel_x_range": [0.5, 0.5],  # X方向の線速度範囲
        "lin_vel_y_range": [0, 0],  # Y方向の線速度範囲
        "ang_vel_range": [0, 0],  # 角速度範囲
    }

    return env_cfg, obs_cfg, reward_cfg, command_cfg

# メイン
def main():
    # コマンドライン引数の解析
    parser = argparse.ArgumentParser()
    parser.add_argument("-e", "--exp_name", type=str, default="go2-walking")
    parser.add_argument("-B", "--num_envs", type=int, default=4096)
    parser.add_argument("--max_iterations", type=int, default=100)
    args = parser.parse_args()

    # Genesisの初期化
    gs.init(logging_level="warning")

    # 各種設定の取得
    log_dir = f"logs/{args.exp_name}"
    env_cfg, obs_cfg, reward_cfg, command_cfg = get_cfgs()
    train_cfg = get_train_cfg(args.exp_name, args.max_iterations)

    # ログディレクトリが存在する場合は削除し再作成
    if os.path.exists(log_dir):
        shutil.rmtree(log_dir)
    os.makedirs(log_dir, exist_ok=True)

    # 環境の初期化
    env = Go2Env(
        num_envs=args.num_envs, 
        env_cfg=env_cfg, 
        obs_cfg=obs_cfg, 
        reward_cfg=reward_cfg, 
        command_cfg=command_cfg
    )

    # ランナーの初期化
    runner = OnPolicyRunner(env, train_cfg, log_dir, device="cuda:0")

    # 設定の保存 (pickle)
    pickle.dump(
        [env_cfg, obs_cfg, reward_cfg, command_cfg, train_cfg],
        open(f"{log_dir}/cfgs.pkl", "wb"),
    )

    # 学習の開始
    runner.learn(
        num_learning_iterations=args.max_iterations, 
        init_at_random_ep_len=True
    )

if __name__ == "__main__":
    main()

4. go2_eval.py

強化学習環境 (go2_env.py) を使用して評価 (動作確認) するコードです。

コードは、次のとおりです。

import argparse
import os
import pickle
import torch
from go2_env import Go2Env
from rsl_rl.runners import OnPolicyRunner
import genesis as gs

# メイン
def main():
    # コマンドライン引数の解析
    parser = argparse.ArgumentParser()
    parser.add_argument("-e", "--exp_name", type=str, default="go2-walking")
    parser.add_argument("--ckpt", type=int, default=100)
    args = parser.parse_args()

    # Genesisの初期化
    gs.init()

    # 各種設定の取得
    log_dir = f"logs/{args.exp_name}"
    env_cfg, obs_cfg, reward_cfg, command_cfg, train_cfg = pickle.load(open(f"logs/{args.exp_name}/cfgs.pkl", "rb"))
    reward_cfg["reward_scales"] = {}

    # 環境の初期化
    env = Go2Env(
        num_envs=1,  # 環境数を指定
        env_cfg=env_cfg,  # 環境設定
        obs_cfg=obs_cfg,  # 観測設定
        reward_cfg=reward_cfg,  # 報酬設定
        command_cfg=command_cfg,  # 命令設定
        show_viewer=True,  # ビューアを表示するかどうか
    )

    # OnPolicyRunnerの初期化
    runner = OnPolicyRunner(env, train_cfg, log_dir, device="cuda:0")

    # モデルの読み込み
    resume_path = os.path.join(log_dir, f"model_{args.ckpt}.pt")
    runner.load(resume_path)
    policy = runner.get_inference_policy(device="cuda:0")

    # 環境のリセット
    obs, _ = env.reset()

    # 推論ループ
    with torch.no_grad():
        while True:
            actions = policy(obs)  # ポリシーで行動を計算
            obs, _, rews, dones, infos = env.step(actions)  # 環境を1ステップ進行


if __name__ == "__main__":
    main()

次回