この記事は強化学習 Advent Calendar 2021の12/22の記事です。
はじめまして！東京大学の吉田です。
大学では身体を持って自律的に発達する人工知能をつくることに興味があって、それを研究しています。
今回の記事は強化学習というよりは強化学習の環境をmujoco-pyを使って作るときのtipsといった内容です。 すでに強化学習の環境を作ってみたり、mujoco-pyを使ってMujocoをつかった物理シミュレーションをしている・やろうとしている人向けの内容です。
ここで紹介するサンプルは全て以下のレポジトリから利用できます。 https://github.com/ugo-nama-kun/mujoco_marker_example
MujocoはEmo Todorovらによって開発された物理シミュレータで、これまでの物理シミュレータでは困難であった物体同士の接触を高速かつより正確に扱うことを可能としています。
DeepMindがMujocoの権利を取得し、今や誰でもフリーに使うことができるようになったのは強化学習界隈で2021年のちょっとしたニュースでした。ライセンス関係は確かに面倒だったので、良い時代になったのものです。ソフトウェアのソースがオープンになれば、ゲームやコンピュータグラフィクスに新しい表現を可能とするのではないかと密かに期待しています。
いっぽう、mujoco-pyはOpenAIが提供しているMujocoをpythonで扱えるよう作成されたライブラリです。v2.1.2.14に更新されてからオープンになったMujocoに対応されてインストールが簡単になりました :) 残念ながら公式ドキュメントはあまり整備されているわけではないため、正直今でもかなり試行錯誤が必要な時があります笑
今回はmujoco-pyが提供する"marker"という便利機能について紹介します。
add_markerはmujoco-pyの公式レポジトリのexamplesにひっそりと紹介されている便利機能です。
公式のドキュメントはありません。。
このadd_markerはシミュレータの任意の座標に、物理的な接触をしない物体を表示することができる機能です。
mujoco-pyを直に触ったことがない人でも、OpenAI Gymは触ったこともある人は多いかもしれません。とりあえず強化学習してみるときに便利ですよね！OpenAI GymのMujocoタスクでは、mujoco-pyが使われています。
でも、mujocoを使っていると環境の定義がどうなっているのかよくわからない時があります。 例えば、次の例はSwimmerというタスクの環境のステップの定義です
class SwimmerEnv(mujoco_env.MujocoEnv, utils.EzPickle):
...
def step(self, a):
ctrl_cost_coeff = 0.0001
xposbefore = self.sim.data.qpos[0]
self.do_simulation(a, self.frame_skip)
xposafter = self.sim.data.qpos[0]
reward_fwd = (xposafter - xposbefore) / self.dt
reward_ctrl = -ctrl_cost_coeff * np.square(a).sum()
reward = reward_fwd + reward_ctrl
ob = self._get_obs()
return ob, reward, False, dict(reward_fwd=reward_fwd, reward_ctrl=reward_ctrl)
...こういったとき、「self.sim.data.qpos[0]ってなんやねん」と思いませんか？私は思いました。
こんなとき、とりあえずその場所を表示してみるのがいい気がします。そこでadd_markerを使うとこんなふうに赤いボールで可視化できます↓
これはけっこうやばそうなところにあるようですね。。本当はまんなかの棒の中央にボールが欲しいところです。実はSwimmerの学習では変な局所最適にハマることが知られています。それはこういう変な報酬設定に理由があるのかもしれません。
今の可視化は元のSwimmerEnvを継承した以下のコードで見ることができます（swimmer_marker.py）。
import numpy as np
from gym.envs.mujoco import SwimmerEnv
from mujoco_py.generated import const
class SwimmerMarkerEnv(SwimmerEnv):
def render(self, **kwargs):
if self.viewer:
x, y = self.sim.data.qpos[:2]
# Draw a sphere marker
self.viewer.add_marker(pos=np.array([x, y, 0]), # Position
label=" ", # Text beside the marker
type=const.GEOM_SPHERE, # Geomety type
size=(0.1, 0.1, 0.1), # Size of the marker
rgba=(1, 0, 0, 1)) # RGBA of the marker
# Default swimmer renderer method
super(SwimmerMarkerEnv, self).render()
if __name__ == '__main__':
# Run
env = SwimmerMarkerEnv()
while True:
env.step(env.action_space.sample())
env.render()簡単そうでしょ？
MujocoEnvを継承して作られたgymの環境はself.viewerからadd_markerを使えます。
まずは公式のmujoco-pyでmarkers_demo.pyで紹介されているものをほぼそのままやります。レポジトリのbasic.pyを実行すると、下のようにboxがぐるぐる回るマーカーが表示されます。
以下がコードの主な部分で、viewerのadd_markerでマーカーを追加することができます。このマーカーはviewerのrenderが呼ばれるごとに削除されます。
model = load_model_from_xml(MODEL_XML)
sim = MjSim(model)
viewer = MjViewer(sim)
step = 0
while True:
t = time.time()
x, y = math.cos(t), math.sin(t)
viewer.add_marker(type=const.GEOM_BOX,
pos=np.array([x, y, 1]),
label=str(t))
viewer.render()
step += 1
if step > 100 and os.getenv('TESTING') is not None:
breakもちろん表示できるマーカーは一つだけではなく、1000個まで表示できます（multi.py）。
上の例はBoxでしたが、他にもいろいろな形状のものを表示することができます。 指定できる形状は_mjtGeomで定義されている以下の形状です。
# _mjtGeom
GEOM_PLANE = 0
GEOM_HFIELD = 1
GEOM_SPHERE = 2
GEOM_CAPSULE = 3
GEOM_ELLIPSOID = 4
GEOM_CYLINDER = 5
GEOM_BOX = 6
GEOM_MESH = 7
GEOM_ARROW = 100
GEOM_ARROW1 = 101
GEOM_ARROW2 = 102
GEOM_LINE = 103
GEOM_SKIN = 104
GEOM_LABEL = 105
GEOM_NONE = 1001
これらはmujoco-pyではfrom mujoco_py.generated import constとすることでconst.GEOM_BOXのように取ってこれます。
例えば、上のbasic.pyのBoxをSphereにしたければconst.GEOM_SPHEREとすることで以下のように表示することができます（basic_sphere.py）。
viewer.add_marker(type=const.GEOM_SPHERE,
pos=np.array([x, y, 1]),
label=str(t))
そのほかの特性は_mjvGeomの構造体で定義されている情報をadd_markerで与えてやることで決めることができます（property.py）。
viewer.add_marker(type=type,
pos=np.array([x, y, 1]),
label=" ",
size=size,
rgba=rgba,
emission=emission)
BOXやSPHEREは位置を与えるだけでも表示できますが、ARROWは方向を与えてやる必要があります（arrow.py）。
そのほかLINEやCAPSULEなど色々指定できるようなのですが、使用例もなかなか見つけられずよくわかっていません。。 見つけた方はぜひ紹介記事を書いていただけるとうれしいです :)
3Dモデルがあれば、それをmujocoで読み込んでそれをマーカーに使うこともできます。ここでは試しにバナナを表示させてみましょう（banana.py）。
xml_path = 'banana.xml'
model = load_model_from_path(xml_path)
sim = MjSim(model)
viewer = MjViewer(sim)
step = 0
while True:
t = time.time()
x, y = math.cos(t), math.sin(t)
viewer.add_marker(pos=np.array([np.cos(t), np.sin(t), 1.0]),
type=const.GEOM_MESH,
rgba=(1, 1, 0, 1),
dataid=0)
viewer.render()ここでbanana.xmlは以下のxmlファイルです。重要なのはassetでbanana.stlを読み込んでいることで、ここで読み込んだデータをadd_markerのdataid=0で指定してメッシュとして使っています。MujocoはSTLをメッシュのアセットとして使えるので、適当な3Dファイルを例えばTurbosquidなどで拾ってきてMeshlabでいい感じに加工することでMujocoで利用できるようになります。
<?xml version="1.0" ?>
<mujoco>
<asset>
<mesh file="banana.stl" name="banana"></mesh>
</asset>
<worldbody>
<body name="box" pos="0 0 0.2">
<geom size="0.15 0.15 0.15" type="box"/>
<joint axis="1 0 0" name="box:x" type="slide"/>
<joint axis="0 1 0" name="box:y" type="slide"/>
</body>
<body name="floor" pos="0 0 0.025">
<geom size="1.0 1.0 0.02" rgba="0 1 0 1" type="box"/>
</body>
</worldbody>
</mujoco>これを実行すると、このように回るバナナを表示できます。
一番簡単なのは、使用例で示したようにGEOM_BOXやGEOM_SPHEREを使ってrenderメソッドに表示させるものを追加していくことです。1000個まで描画可能なので、これでいろいろなものを表示することができます（swimmer_party.py）。
xmlでassetを読んでこなければいけないため、そのせいでいろいろと面倒ですがgymの環境でもバナナを表示させることができます（banana_gym.py）。もっと楽にする方法をご存知でしたら教えてください。。
実はmujoco-pyはrgb_arrayを出力する際にマーカーをクリアしてくれません（openai/mujoco-py#423）。なので、そのまま実行すると以下のエラーが出ます。
RuntimeError: Ran out of geoms. maxgeom: 1000
これを解決するにはdel viewer._markers[:]をrender()を読んだ後に呼ぶ必要があります。
※ 正確には、MjRenderContextOffscreen(sim, -1)がmujocoのviewerでセットされている時です。gymではrgb_arrayを出すときにセットされています。
ちょっと何かを表示させたいときやデバッグに使えるmujoco_pyのマーカー機能でした。
みなさまのMujocoライフがより捗る手助けになれば幸いです。
メリークリスマース🎅🎄🎁