Skip to content

Instantly share code, notes, and snippets.

@ugo-nama-kun
Last active May 30, 2022
Embed
What would you like to do?
mujoco_py の marker の使い方

mujoco_py のマーカーの使い方

この記事は強化学習 Advent Calendar 2021の12/22の記事です。

はじめまして!東京大学の吉田です。

大学では身体を持って自律的に発達する人工知能をつくることに興味があって、それを研究しています。

今回の記事は強化学習というよりは強化学習の環境をmujoco-pyを使って作るときのtipsといった内容です。 すでに強化学習の環境を作ってみたり、mujoco-pyを使ってMujocoをつかった物理シミュレーションをしている・やろうとしている人向けの内容です。

ここで紹介するサンプルは全て以下のレポジトリから利用できます。 https://github.com/ugo-nama-kun/mujoco_marker_example

Mujocoって? mujoco-pyって?

MujocoはEmo Todorovらによって開発された物理シミュレータで、これまでの物理シミュレータでは困難であった物体同士の接触を高速かつより正確に扱うことを可能としています。

DeepMindがMujocoの権利を取得し、今や誰でもフリーに使うことができるようになったのは強化学習界隈で2021年のちょっとしたニュースでした。ライセンス関係は確かに面倒だったので、良い時代になったのものです。ソフトウェアのソースがオープンになれば、ゲームやコンピュータグラフィクスに新しい表現を可能とするのではないかと密かに期待しています。

いっぽう、mujoco-pyはOpenAIが提供しているMujocoをpythonで扱えるよう作成されたライブラリです。v2.1.2.14に更新されてからオープンになったMujocoに対応されてインストールが簡単になりました :) 残念ながら公式ドキュメントはあまり整備されているわけではないため、正直今でもかなり試行錯誤が必要な時があります笑

今回はmujoco-pyが提供する"marker"という便利機能について紹介します。

"Marker" って?

add_markerはmujoco-pyの公式レポジトリのexamplesにひっそりと紹介されている便利機能です。 公式のドキュメントはありません。。

このadd_markerはシミュレータの任意の座標に、物理的な接触をしない物体を表示することができる機能です。

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_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を使えます。

Markerの使い方

基本

まずは公式のmujoco-pyでmarkers_demo.pyで紹介されているものをほぼそのままやります。レポジトリのbasic.pyを実行すると、下のようにboxがぐるぐる回るマーカーが表示されます。

basic

以下がコードの主な部分で、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)。

multi

他の形状を表示する

上の例は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))

sphere

そのほかの特性は_mjvGeomの構造体で定義されている情報をadd_markerで与えてやることで決めることができます(property.py)。

viewer.add_marker(type=type,
                  pos=np.array([x, y, 1]),
                  label=" ",
                  size=size,
                  rgba=rgba,
                  emission=emission)

prop

BOXやSPHEREは位置を与えるだけでも表示できますが、ARROWは方向を与えてやる必要があります(arrow.py)。

arrow

そのほかLINEやCAPSULEなど色々指定できるようなのですが、使用例もなかなか見つけられずよくわかっていません。。 見つけた方はぜひ紹介記事を書いていただけるとうれしいです :)

Meshを表示する

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_markerdataid=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>

これを実行すると、このように回るバナナを表示できます。

banana

Gymで使う

一番簡単なのは、使用例で示したようにGEOM_BOXやGEOM_SPHEREを使ってrenderメソッドに表示させるものを追加していくことです。1000個まで描画可能なので、これでいろいろなものを表示することができます(swimmer_party.py)。

swimmer_party

xmlでassetを読んでこなければいけないため、そのせいでいろいろと面倒ですがgymの環境でもバナナを表示させることができます(banana_gym.py)。もっと楽にする方法をご存知でしたら教えてください。。

banana_ant

rgb_arrayを取得する時に気をつけること

実は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ライフがより捗る手助けになれば幸いです。

メリークリスマース🎅🎄🎁

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment