0. ある日のこと

「あれ？転倒してるのにエピソード終了しないんだ．．．」

MuJoCoのAnt-v2環境においてエピソードが途中で終了してしまう条件は以下のようになっています．

# 終了条件，三つのうちどれか一つでもFalseだとエピソードが終了．
notdone = np.isfinite(state).all() and state[2] >= 0.2 and state[2] <= 1.0
done = not notdone

Ant-v2環境では，転倒して歩行していない状態でもエピソードが終了するまでsurvive rewardとして報酬を獲得することができます．一方エピソードの終了条件には，エージェントの高さが一定以上になると終了するというものがあります．これにより積極的に歩行してジャンプしてしまったエージェントはエピソードがそこで終了するため，転倒してsurvive rewardを稼いでいるエージェントより報酬が低くなってしまうことがあります．学習結果を見ているときにこの仕様はいかがなものかと思い，MuJoCoのリファレンス及び，内部の仕様について調査して分かったことがあるのでまとめます．

この記事では，強化学習界隈でよく使われるMuJoCoで，エージェントの体の向きを取得する方法について説明します．

この度，ボディの方向ベクトルを抽出するプログラムをパッケージ化しました．
「方向ベクトル取得したかった！」という方はこれをインストールするだけで利用できると思います．
MuJoCoのリファレンスは癖が強いので，今回の実装はかなり苦労しました．

1. クォータニオン（Quaternion）とは

MuJoCoでは，クォータニオン(Quaternion)という形式でロボットの各部位の「回転」, 「姿勢」を取得することができます．

クォータニオンは要素数4の回転行列のようなものです．計算も回転行列の要領で行うことができます．
より詳しい解説は以下の記事などが参考になります．

クォータニオンは，シミュレーションや3Dゲームなどでよく使われている数学概念です．
馴染みの深い回転行列やオイラー角などはパラメータ数の多さ，計算時間的問題があるため，このクォータニオンが盛んに使われているらしいです．

2. MuJoCoでの実装

私の場合やりたいこととしては，エージェントの転倒を検出することです．これはつまりエージェントの状態（向き，姿勢）を取得することです．これを実現するためにはエージェントのクォータニオンを取得して，胴体がどの方向を向いているのかを算出すれば良いと分かります．

リファレンスでそれっぽい関数を見つけました．

void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]);

こちら，ベクトルをクォータニオンで回転させるという機能を持っています．使い方としては，第一引数に計算結果を格納する要素数3の配列，第二引数に基準となるベクトル，第三引数にボディのクォータニオンを渡すだけです．
次節に各引数について解説します．

2-1. ボディのクォータニオン

現時点のロボットのクォータニオンは以下の関数で取得できます．
引数nameでクォータニオンを取得したい部位を指定できます．

get_body_xquat(name)

MuJoCoなど多くのシミュレーション環境では，ロボットをxmlファイルで定義しています．例えばMuJoCoのAnt-v2環境では，antロボットを以下のように定義しています．

ant.xml

<mujoco model="ant">
  <compiler angle="degree" coordinate="local" inertiafromgeom="true"/>
  <option integrator="RK4" timestep="0.01"/>
  <custom>
    <numeric data="0.0 0.0 0.55 1.0 0.0 0.0 0.0 0.0 1.0 0.0 -1.0 0.0 -1.0 0.0 1.0" name="init_qpos"/>
  </custom>
  <default>
    <joint armature="1" damping="1" limited="true"/>
    <geom conaffinity="0" condim="3" density="5.0" friction="1 0.5 0.5" margin="0.01" rgba="0.8 0.6 0.4 1"/>
  </default>
  <asset>
    <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>
    <texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
    <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
    <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
    <material name="geom" texture="texgeom" texuniform="true"/>
  </asset>
  <worldbody>
    <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
    <geom conaffinity="1" condim="3" material="MatPlane" name="floor" pos="0 0 0" rgba="0.8 0.9 0.8 1" size="40 40 40" type="plane"/>
    <body name="torso" pos="0 0 0.75">
      <camera name="track" mode="trackcom" pos="0 -3 0.3" xyaxes="1 0 0 0 0 1"/>
      <geom name="torso_geom" pos="0 0 0" size="0.25" type="sphere"/>
      <joint armature="0" damping="0" limited="false" margin="0.01" name="root" pos="0 0 0" type="free"/>
      <body name="front_left_leg" pos="0 0 0">
        <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="aux_1_geom" size="0.08" type="capsule"/>
        <body name="aux_1" pos="0.2 0.2 0">
          <joint axis="0 0 1" name="hip_1" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
          <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="left_leg_geom" size="0.08" type="capsule"/>
          <body pos="0.2 0.2 0">
            <joint axis="-1 1 0" name="ankle_1" pos="0.0 0.0 0.0" range="30 70" type="hinge"/>
            <geom fromto="0.0 0.0 0.0 0.4 0.4 0.0" name="left_ankle_geom" size="0.08" type="capsule"/>
          </body>
        </body>
      </body>
      <body name="front_right_leg" pos="0 0 0">
        <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="aux_2_geom" size="0.08" type="capsule"/>
        <body name="aux_2" pos="-0.2 0.2 0">
          <joint axis="0 0 1" name="hip_2" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
          <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="right_leg_geom" size="0.08" type="capsule"/>
          <body pos="-0.2 0.2 0">
            <joint axis="1 1 0" name="ankle_2" pos="0.0 0.0 0.0" range="-70 -30" type="hinge"/>
            <geom fromto="0.0 0.0 0.0 -0.4 0.4 0.0" name="right_ankle_geom" size="0.08" type="capsule"/>
          </body>
        </body>
      </body>
      <body name="back_leg" pos="0 0 0">
        <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="aux_3_geom" size="0.08" type="capsule"/>
        <body name="aux_3" pos="-0.2 -0.2 0">
          <joint axis="0 0 1" name="hip_3" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
          <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="back_leg_geom" size="0.08" type="capsule"/>
          <body pos="-0.2 -0.2 0">
            <joint axis="-1 1 0" name="ankle_3" pos="0.0 0.0 0.0" range="-70 -30" type="hinge"/>
            <geom fromto="0.0 0.0 0.0 -0.4 -0.4 0.0" name="third_ankle_geom" size="0.08" type="capsule"/>
          </body>
        </body>
      </body>
      <body name="right_back_leg" pos="0 0 0">
        <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="aux_4_geom" size="0.08" type="capsule"/>
        <body name="aux_4" pos="0.2 -0.2 0">
          <joint axis="0 0 1" name="hip_4" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
          <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="rightback_leg_geom" size="0.08" type="capsule"/>
          <body pos="0.2 -0.2 0">
            <joint axis="1 1 0" name="ankle_4" pos="0.0 0.0 0.0" range="30 70" type="hinge"/>
            <geom fromto="0.0 0.0 0.0 0.4 -0.4 0.0" name="fourth_ankle_geom" size="0.08" type="capsule"/>
          </body>
        </body>
      </body>
    </body>
  </worldbody>
  <actuator>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_4" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_4" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_1" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_1" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_2" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_2" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_3" gear="150"/>
    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_3" gear="150"/>
  </actuator>
</mujoco>

この中の <body> タグのname要素がエージェントの部位の名前になります．

2-2. 基準となるベクトル

ここで指定するベクトルを現在のクォータニオンで回転させることで，基準となるベクトルが現在どの方向を向いているのかが分かります．
つまり，エージェントの転倒を検出したい時，z方向のベクトル(0,0,1)を基準とし，戻り値のベクトルのz値が-1付近であるかどうかを確認すれば良いのです．

2-3. rotVecQuat関数の使い方

クォータニオンは1ステップごとに更新されるので，この関数を利用するとしたら，gym環境をラップしてstep()関数に追記するのが良いでしょう．
以下に例を示します．

import numpy as np 
import gym
import mujoco_py

# ボディの方向ベクトルを取得する自作環境
class GetOrientationEnv(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        # 基準となるベクトル
        self.base_vec = np.array([0.0, 0.0, 1.0])

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        # 胴体のクォータニオンの取得
        quat = self.data.get_body_xquat('torso')
        # 計算結果を格納する変数の用意
        body_orientation = np.zeros(3) 
        # 計算
        mujoco_py.functions.mju_rotVecQuat(body_orientation, self.base_vec, quat)
        info['body_orientation'] = body_orientation
        return obs, reward, done, info

def main():
    env = gym.make('Ant-v2')
    # 自作環境でラップ
    env = GetOrientationEnv(env)

    done = False
    env.reset()
    while True:
        env.render()
        action = env.action_space.sample()
        obs, reward, done, info = env.step(action)
        # 方向ベクトルが返ってくる！
        print(info['body_orientation'])
        if done:
            env.reset()
            break

if __name__=='__main__':
    main()

これで晴れて，エージェントの転倒を検出することができます！
冒頭にも載せたこのgifは取得した方向ベクトルを可視化したものです．

しっかり取得できていることがわかりますね．

3. おわりに

MuJoCoを使いこなすのは大変だなぁ．．．