Repository: HumanCompatibleAI/imitation (language: python)

Issue:
Relocate / delete example environments
https://github.com/HumanCompatibleAI/imitation/pull/55 imported a bunch of environments originally used in the AIRL paper into `imitation`. This made sense at the time since we wanted to benchmark them, and we didn't have any obvious place to put them.

However, we're now as far as I can tell mostly not using them for benchmarking purposes (using more standard Gym environments instead), and we also have `seals` as a place to centralize reward/imitation-learning relevant environments.

These environments now feel a bit out of place in `imitation`. Because of this I'm inclined to relocate the environments we actually want to use to `seals` and delete the rest.

@kmdanielduan's project may give us good alternatives to TwoDMaze and also to DisabledAnt. And I think we don't really need ObjPusher or CustomAnt?


Discussion / hints:

Fix patch:
diff --git a/experiments/imit_benchmark_config.csv b/experiments/imit_benchmark_config.csv
index 4605ea6a7..3da073ab1 100644
--- a/experiments/imit_benchmark_config.csv
+++ b/experiments/imit_benchmark_config.csv
@@ -38,15 +38,3 @@ reacher,50000,4
 reacher,50000,11
 reacher,50000,18
 reacher,50000,25
-two_d_maze,50000,4
-two_d_maze,50000,11
-two_d_maze,50000,18
-two_d_maze,50000,25
-custom_ant,50000,4
-custom_ant,50000,11
-custom_ant,50000,18
-custom_ant,50000,25
-disabled_ant,50000,4
-disabled_ant,50000,11
-disabled_ant,50000,18
-disabled_ant,50000,25
diff --git a/experiments/rollouts_from_policies_config.csv b/experiments/rollouts_from_policies_config.csv
index 92d3d6039..b1ef9eb50 100644
--- a/experiments/rollouts_from_policies_config.csv
+++ b/experiments/rollouts_from_policies_config.csv
@@ -11,6 +11,3 @@ swimmer,40
 ant,40
 humanoid,240
 reacher,40
-two_d_maze,40
-custom_ant,40
-disabled_ant,40
diff --git a/experiments/train_experts.sh b/experiments/train_experts.sh
index f2a19205e..953206697 100755
--- a/experiments/train_experts.sh
+++ b/experiments/train_experts.sh
@@ -7,7 +7,6 @@ set -e
 
 ENVS+="acrobot cartpole mountain_car "
 ENVS+="reacher half_cheetah hopper ant humanoid swimmer walker "
-ENVS+="two_d_maze custom_ant disabled_ant "
 
 SEEDS="0 1 2"
 
diff --git a/src/imitation/envs/examples/__init__.py b/src/imitation/envs/examples/__init__.py
index e836e5c46..162d0d26a 100644
--- a/src/imitation/envs/examples/__init__.py
+++ b/src/imitation/envs/examples/__init__.py
@@ -5,4 +5,4 @@
 """
 
 # Register environments with Gym
-from imitation.envs.examples import airl_envs, model_envs  # noqa: F401
+from imitation.envs.examples import model_envs  # noqa: F401
diff --git a/src/imitation/envs/examples/airl_envs/LICENSE b/src/imitation/envs/examples/airl_envs/LICENSE
deleted file mode 100644
index 6ec89f715..000000000
--- a/src/imitation/envs/examples/airl_envs/LICENSE
+++ /dev/null
@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2018 Justin Fu
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
diff --git a/src/imitation/envs/examples/airl_envs/__init__.py b/src/imitation/envs/examples/airl_envs/__init__.py
deleted file mode 100644
index a8053e734..000000000
--- a/src/imitation/envs/examples/airl_envs/__init__.py
+++ /dev/null
@@ -1,42 +0,0 @@
-from typing import Optional
-
-from gym.envs import register as gym_register
-
-_ENTRY_POINT_PREFIX = "imitation.envs.examples.airl_envs"
-
-
-def _register(env_name: str, entry_point: str, kwargs: Optional[dict] = None):
-    entry_point = f"{_ENTRY_POINT_PREFIX}.{entry_point}"
-    gym_register(id=env_name, entry_point=entry_point, kwargs=kwargs)
-
-
-def _point_maze_register():
-    for dname, dval in {"Left": 0, "Right": 1}.items():
-        for vname, vval in {"": False, "Vel": True}.items():
-            _register(
-                f"imitation/PointMaze{dname}{vname}-v0",
-                entry_point="point_maze_env:PointMazeEnv",
-                kwargs={"direction": dval, "include_vel": vval},
-            )
-
-
-_register(
-    "imitation/ObjPusher-v0",
-    entry_point="pusher_env:PusherEnv",
-    kwargs={"sparse_reward": False},
-)
-_register("imitation/TwoDMaze-v0", entry_point="twod_maze:TwoDMaze")
-
-_point_maze_register()
-
-# A modified ant which flips over less and learns faster via TRPO
-_register(
-    "imitation/CustomAnt-v0",
-    entry_point="ant_env:CustomAntEnv",
-    kwargs={"gear": 30, "disabled": False},
-)
-_register(
-    "imitation/DisabledAnt-v0",
-    entry_point="ant_env:CustomAntEnv",
-    kwargs={"gear": 30, "disabled": True},
-)
diff --git a/src/imitation/envs/examples/airl_envs/ant_env.py b/src/imitation/envs/examples/airl_envs/ant_env.py
deleted file mode 100644
index f70e8c786..000000000
--- a/src/imitation/envs/examples/airl_envs/ant_env.py
+++ /dev/null
@@ -1,718 +0,0 @@
-import numpy as np
-from gym import utils
-from gym.envs.mujoco import mujoco_env
-
-from imitation.envs.examples.airl_envs.dynamic_mjc.model_builder import MJCModel
-
-
-def ant_env(gear=150, eyes=True):
-    mjcmodel = MJCModel("ant_maze")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
-    mjcmodel.root.option(timestep="0.01", integrator="RK4")
-    mjcmodel.root.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",
-    )
-    asset = mjcmodel.root.asset()
-    asset.texture(
-        builtin="gradient",
-        height="100",
-        rgb1="1 1 1",
-        rgb2="0 0 0",
-        type="skybox",
-        width="100",
-    )
-    asset.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",
-    )
-    asset.texture(
-        builtin="checker",
-        height="100",
-        name="texplane",
-        rgb1="0 0 0",
-        rgb2="0.8 0.8 0.8",
-        type="2d",
-        width="100",
-    )
-    asset.material(
-        name="MatPlane",
-        reflectance="0.5",
-        shininess="1",
-        specular="1",
-        texrepeat="60 60",
-        texture="texplane",
-    )
-    asset.material(name="geom", texture="texgeom", texuniform="true")
-
-    default = mjcmodel.root.default()
-    default.joint(armature=1, damping=1, limited="true")
-    default.geom(
-        friction=[1.5, 0.5, 0.5],
-        density=5.0,
-        margin=0.01,
-        condim=3,
-        conaffinity=0,
-        rgba="0.8 0.6 0.4 1",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-    worldbody.light(
-        cutoff="100",
-        diffuse=[0.8, 0.8, 0.8],
-        dir="-0 0 -1.3",
-        directional="true",
-        exponent="1",
-        pos="0 0 1.3",
-        specular=".1 .1 .1",
-    )
-    worldbody.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",
-    )
-
-    ant = worldbody.body(name="torso", pos=[0, 0, 0.75])
-    ant.geom(name="torso_geom", pos=[0, 0, 0], size="0.25", type="sphere")
-    ant.joint(
-        armature="0",
-        damping="0",
-        limited="false",
-        margin="0.01",
-        name="root",
-        pos=[0, 0, 0],
-        type="free",
-    )
-
-    if eyes:
-        eye_z = 0.1
-        eye_y = -0.21
-        eye_x_offset = 0.07
-        # eyes
-        ant.geom(
-            fromto=[eye_x_offset, 0, eye_z, eye_x_offset, eye_y, eye_z],
-            name="eye1",
-            size="0.03",
-            type="capsule",
-            rgba=[1, 1, 1, 1],
-        )
-        ant.geom(
-            fromto=[eye_x_offset, 0, eye_z, eye_x_offset, eye_y - 0.02, eye_z],
-            name="eye1_",
-            size="0.02",
-            type="capsule",
-            rgba=[0, 0, 0, 1],
-        )
-        ant.geom(
-            fromto=[-eye_x_offset, 0, eye_z, -eye_x_offset, eye_y, eye_z],
-            name="eye2",
-            size="0.03",
-            type="capsule",
-            rgba=[1, 1, 1, 1],
-        )
-        ant.geom(
-            fromto=[-eye_x_offset, 0, eye_z, -eye_x_offset, eye_y - 0.02, eye_z],
-            name="eye2_",
-            size="0.02",
-            type="capsule",
-            rgba=[0, 0, 0, 1],
-        )
-        # eyebrows
-        ant.geom(
-            fromto=[
-                eye_x_offset - 0.03,
-                eye_y,
-                eye_z + 0.07,
-                eye_x_offset + 0.03,
-                eye_y,
-                eye_z + 0.1,
-            ],
-            name="brow1",
-            size="0.02",
-            type="capsule",
-            rgba=[0, 0, 0, 1],
-        )
-        ant.geom(
-            fromto=[
-                -eye_x_offset + 0.03,
-                eye_y,
-                eye_z + 0.07,
-                -eye_x_offset - 0.03,
-                eye_y,
-                eye_z + 0.1,
-            ],
-            name="brow2",
-            size="0.02",
-            type="capsule",
-            rgba=[0, 0, 0, 1],
-        )
-
-    front_left_leg = ant.body(name="front_left_leg", pos=[0, 0, 0])
-    front_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="aux_1_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_1 = front_left_leg.body(name="aux_1", pos=[0.2, 0.2, 0])
-    aux_1.joint(
-        axis=[0, 0, 1],
-        name="hip_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="left_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_1 = aux_1.body(pos=[0.2, 0.2, 0])
-    ankle_1.joint(
-        axis=[-1, 1, 0],
-        name="ankle_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, 0.4, 0.0],
-        name="left_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    front_right_leg = ant.body(name="front_right_leg", pos=[0, 0, 0])
-    front_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="aux_2_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_2 = front_right_leg.body(name="aux_2", pos=[-0.2, 0.2, 0])
-    aux_2.joint(
-        axis=[0, 0, 1],
-        name="hip_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="right_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_2 = aux_2.body(pos=[-0.2, 0.2, 0])
-    ankle_2.joint(
-        axis=[1, 1, 0],
-        name="ankle_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, 0.4, 0.0],
-        name="right_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    back_left_leg = ant.body(name="back_left_leg", pos=[0, 0, 0])
-    back_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="aux_3_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_3 = back_left_leg.body(name="aux_3", pos=[-0.2, -0.2, 0])
-    aux_3.joint(
-        axis=[0, 0, 1],
-        name="hip_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_3.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="backleft_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_3 = aux_3.body(pos=[-0.2, -0.2, 0])
-    ankle_3.joint(
-        axis=[-1, 1, 0],
-        name="ankle_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_3.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, -0.4, 0.0],
-        name="backleft_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    back_right_leg = ant.body(name="back_right_leg", pos=[0, 0, 0])
-    back_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="aux_4_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_4 = back_right_leg.body(name="aux_4", pos=[0.2, -0.2, 0])
-    aux_4.joint(
-        axis=[0, 0, 1],
-        name="hip_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_4.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="backright_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_4 = aux_4.body(pos=[0.2, -0.2, 0])
-    ankle_4.joint(
-        axis=[1, 1, 0],
-        name="ankle_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_4.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, -0.4, 0.0],
-        name="backright_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_4", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_4", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_1", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_1", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_2", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_2", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_3", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_3", gear=gear)
-    return mjcmodel
-
-
-def angry_ant_crippled(gear=150):
-    mjcmodel = MJCModel("ant_maze")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
-    mjcmodel.root.option(timestep="0.01", integrator="RK4")
-    mjcmodel.root.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",
-    )
-    asset = mjcmodel.root.asset()
-    asset.texture(
-        builtin="gradient",
-        height="100",
-        rgb1="1 1 1",
-        rgb2="0 0 0",
-        type="skybox",
-        width="100",
-    )
-    asset.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",
-    )
-    asset.texture(
-        builtin="checker",
-        height="100",
-        name="texplane",
-        rgb1="0 0 0",
-        rgb2="0.8 0.8 0.8",
-        type="2d",
-        width="100",
-    )
-    asset.material(
-        name="MatPlane",
-        reflectance="0.5",
-        shininess="1",
-        specular="1",
-        texrepeat="60 60",
-        texture="texplane",
-    )
-    asset.material(name="geom", texture="texgeom", texuniform="true")
-
-    default = mjcmodel.root.default()
-    default.joint(armature=1, damping=1, limited="true")
-    default.geom(
-        friction=[1.5, 0.5, 0.5],
-        density=5.0,
-        margin=0.01,
-        condim=3,
-        conaffinity=0,
-        rgba="0.8 0.6 0.4 1",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-
-    worldbody.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",
-    )
-    worldbody.light(
-        cutoff="100",
-        diffuse=[0.8, 0.8, 0.8],
-        dir="-0 0 -1.3",
-        directional="true",
-        exponent="1",
-        pos="0 0 1.3",
-        specular=".1 .1 .1",
-    )
-
-    ant = worldbody.body(name="torso", pos=[0, 0, 0.75])
-    ant.geom(name="torso_geom", pos=[0, 0, 0], size="0.25", type="sphere")
-    ant.joint(
-        armature="0",
-        damping="0",
-        limited="false",
-        margin="0.01",
-        name="root",
-        pos=[0, 0, 0],
-        type="free",
-    )
-
-    eye_z = 0.1
-    eye_y = -0.21
-    eye_x_offset = 0.07
-    # eyes
-    ant.geom(
-        fromto=[eye_x_offset, 0, eye_z, eye_x_offset, eye_y, eye_z],
-        name="eye1",
-        size="0.03",
-        type="capsule",
-        rgba=[1, 1, 1, 1],
-    )
-    ant.geom(
-        fromto=[eye_x_offset, 0, eye_z, eye_x_offset, eye_y - 0.02, eye_z],
-        name="eye1_",
-        size="0.02",
-        type="capsule",
-        rgba=[0, 0, 0, 1],
-    )
-    ant.geom(
-        fromto=[-eye_x_offset, 0, eye_z, -eye_x_offset, eye_y, eye_z],
-        name="eye2",
-        size="0.03",
-        type="capsule",
-        rgba=[1, 1, 1, 1],
-    )
-    ant.geom(
-        fromto=[-eye_x_offset, 0, eye_z, -eye_x_offset, eye_y - 0.02, eye_z],
-        name="eye2_",
-        size="0.02",
-        type="capsule",
-        rgba=[0, 0, 0, 1],
-    )
-    # eyebrows
-    ant.geom(
-        fromto=[
-            eye_x_offset - 0.03,
-            eye_y,
-            eye_z + 0.07,
-            eye_x_offset + 0.03,
-            eye_y,
-            eye_z + 0.1,
-        ],
-        name="brow1",
-        size="0.02",
-        type="capsule",
-        rgba=[0, 0, 0, 1],
-    )
-    ant.geom(
-        fromto=[
-            -eye_x_offset + 0.03,
-            eye_y,
-            eye_z + 0.07,
-            -eye_x_offset - 0.03,
-            eye_y,
-            eye_z + 0.1,
-        ],
-        name="brow2",
-        size="0.02",
-        type="capsule",
-        rgba=[0, 0, 0, 1],
-    )
-
-    front_left_leg = ant.body(name="front_left_leg", pos=[0, 0, 0])
-    front_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="aux_1_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_1 = front_left_leg.body(name="aux_1", pos=[0.2, 0.2, 0])
-    aux_1.joint(
-        axis=[0, 0, 1],
-        name="hip_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="left_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_1 = aux_1.body(pos=[0.2, 0.2, 0])
-    ankle_1.joint(
-        axis=[-1, 1, 0],
-        name="ankle_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, 0.4, 0.0],
-        name="left_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    front_right_leg = ant.body(name="front_right_leg", pos=[0, 0, 0])
-    front_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="aux_2_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_2 = front_right_leg.body(name="aux_2", pos=[-0.2, 0.2, 0])
-    aux_2.joint(
-        axis=[0, 0, 1],
-        name="hip_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="right_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_2 = aux_2.body(pos=[-0.2, 0.2, 0])
-    ankle_2.joint(
-        axis=[1, 1, 0],
-        name="ankle_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, 0.4, 0.0],
-        name="right_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    # Back left leg is crippled
-    thigh_length = 0.1  # 0.2
-    ankle_length = 0.2  # 0.4
-    dark_red = [0.8, 0.3, 0.3, 1.0]
-
-    back_left_leg = ant.body(name="back_left_leg", pos=[0, 0, 0])
-    back_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="aux_3_geom",
-        size="0.08",
-        type="capsule",
-        rgba=dark_red,
-    )
-    aux_3 = back_left_leg.body(name="aux_3", pos=[-0.2, -0.2, 0])
-    aux_3.joint(
-        axis=[0, 0, 1],
-        name="hip_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_3.geom(
-        fromto=[0.0, 0.0, 0.0, -thigh_length, -thigh_length, 0.0],
-        name="backleft_leg_geom",
-        size="0.08",
-        type="capsule",
-        rgba=dark_red,
-    )
-    ankle_3 = aux_3.body(pos=[-thigh_length, -thigh_length, 0])
-    ankle_3.joint(
-        axis=[-1, 1, 0],
-        name="ankle_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_3.geom(
-        fromto=[0.0, 0.0, 0.0, -ankle_length, -ankle_length, 0.0],
-        name="backleft_ankle_geom",
-        size="0.08",
-        type="capsule",
-        rgba=dark_red,
-    )
-
-    back_right_leg = ant.body(name="back_right_leg", pos=[0, 0, 0])
-    back_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="aux_4_geom",
-        size="0.08",
-        type="capsule",
-        rgba=dark_red,
-    )
-    aux_4 = back_right_leg.body(name="aux_4", pos=[0.2, -0.2, 0])
-    aux_4.joint(
-        axis=[0, 0, 1],
-        name="hip_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_4.geom(
-        fromto=[0.0, 0.0, 0.0, thigh_length, -thigh_length, 0.0],
-        name="backright_leg_geom",
-        size="0.08",
-        type="capsule",
-        rgba=dark_red,
-    )
-    ankle_4 = aux_4.body(pos=[thigh_length, -thigh_length, 0])
-    ankle_4.joint(
-        axis=[1, 1, 0],
-        name="ankle_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_4.geom(
-        fromto=[0.0, 0.0, 0.0, ankle_length, -ankle_length, 0.0],
-        name="backright_ankle_geom",
-        size="0.08",
-        type="capsule",
-        rgba=dark_red,
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_1", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_1", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_2", gear=gear)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_2", gear=gear)
-    actuator.motor(
-        ctrllimited="true",
-        ctrlrange="-1.0 1.0",
-        joint="hip_3",
-        gear=1,
-    )  # cripple the joints
-    actuator.motor(
-        ctrllimited="true",
-        ctrlrange="-1.0 1.0",
-        joint="ankle_3",
-        gear=1,
-    )  # cripple the joints
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_4", gear=1)
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_4", gear=1)
-    return mjcmodel
-
-
-class CustomAntEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    """A modified ant env with lower joint gear ratios so it flips less often and learns
-    faster."""
-
-    def __init__(self, max_timesteps=1000, disabled=False, gear=150):
-        self.timesteps = 0
-        self.max_timesteps = max_timesteps
-
-        if disabled:
-            model = angry_ant_crippled(gear=gear)
-        else:
-            model = ant_env(gear=gear)
-
-        with model.asfile() as f:
-            mujoco_env.MujocoEnv.__init__(self, f.name, 5)
-
-    def step(self, a):
-        vel = self.sim.data.qvel.flat[0]
-        forward_reward = vel
-        self.do_simulation(a, self.frame_skip)
-
-        ctrl_cost = 0.01 * np.square(a).sum()
-        contact_cost = (
-            0.5 * 1e-3 * np.sum(np.square(np.clip(self.sim.data.cfrc_ext, -1, 1)))
-        )
-        state = self.state_vector()
-        flipped = not (state[2] >= 0.2)
-        flipped_rew = -1 if flipped else 0
-        reward = forward_reward - ctrl_cost - contact_cost + flipped_rew
-
-        self.timesteps += 1
-        done = self.timesteps >= self.max_timesteps
-
-        ob = self._get_obs()
-        return (
-            ob,
-            reward,
-            done,
-            dict(
-                reward_forward=forward_reward,
-                reward_ctrl=-ctrl_cost,
-                reward_contact=-contact_cost,
-                reward_flipped=flipped_rew,
-            ),
-        )
-
-    def _get_obs(self):
-        return np.concatenate(
-            [
-                self.sim.data.qpos.flat[2:],
-                self.sim.data.qvel.flat,
-                np.clip(self.sim.data.cfrc_ext, -1, 1).flat,
-            ],
-        )
-
-    def reset_model(self):
-        self.timesteps = 0
-        qpos = self.init_qpos + self.np_random.uniform(
-            size=self.model.nq,
-            low=-0.1,
-            high=0.1,
-        )
-        qvel = self.init_qvel + self.np_random.randn(self.model.nv) * 0.1
-        self.set_state(qpos, qvel)
-        return self._get_obs()
-
-    def viewer_setup(self):
-        self.viewer.cam.distance = self.model.stat.extent * 0.5
diff --git a/src/imitation/envs/examples/airl_envs/assets/twod_maze.xml b/src/imitation/envs/examples/airl_envs/assets/twod_maze.xml
deleted file mode 100644
index 4b524dea8..000000000
--- a/src/imitation/envs/examples/airl_envs/assets/twod_maze.xml
+++ /dev/null
@@ -1,36 +0,0 @@
-<mujoco model="twod_point">
-    <compiler inertiafromgeom="true" angle="radian" coordinate="local" />
-    <option timestep="0.01" gravity="0 0 0" iterations="20" integrator="Euler" />
-    <default>
-        <joint limited="false" damping="1" />
-        <geom contype="2" conaffinity="1" condim="1" friction=".5 .1 .1" density="1000" margin="0.002" />
-    </default>
-
-    <worldbody>
-        <!-- Pointmass -->
-        <body name="particle" pos="-0.15 -0.15 0">
-            <geom name="particle_geom" type="sphere" size="0.03" rgba="0.0 0.0 1.0 1" contype="1"/>
-            <site name="particle_site" pos="0 0 0" size="0.01" />
-            <joint name="ball_x" type="slide" pos="0 0 0" axis="1 0 0" />
-            <joint name="ball_y" type="slide" pos="0 0 0" axis="0 1 0" />
-        </body>
-
-        <!-- Target -->
-        <body name="target" pos="0.15 -0.15 0">
-            <geom conaffinity="2" name="target_geom" type="sphere" size="0.03" rgba="0 0.9 0.1 1"/>
-        </body>
-
-        <!-- Arena -->
-        <geom conaffinity="1" fromto="-.3 -.3 .01 .3 -.3 .01" name="sideS" rgba="0.9 0.4 0.6 1" size=".02" type="capsule"/>
-		<geom conaffinity="1" fromto=" .3 -.3 .01 .3  .3 .01" name="sideE" rgba="0.9 0.4 0.6 1" size=".02" type="capsule"/>
-		<geom conaffinity="1" fromto="-.3  .3 .01 .3  .3 .01" name="sideN" rgba="0.9 0.4 0.6 1" size=".02" type="capsule"/>
-		<geom conaffinity="1" fromto="-.3 -.3 .01 -.3 .3 .01" name="sideW" rgba="0.9 0.4 0.6 1" size=".02" type="capsule"/>
-
-        <geom conaffinity="1" fromto="0 -.3 .01 0 -.05 .01" name="wall" rgba="0.9 0.4 0.6 1" size=".02" type="capsule"/>
-    </worldbody>
-
-    <actuator>
-        <motor joint="ball_x" ctrlrange="-5.0 5.0" ctrllimited="true"/>
-        <motor joint="ball_y" ctrlrange="-5.0 5.0" ctrllimited="true"/>
-    </actuator>
-</mujoco>
\ No newline at end of file
diff --git a/src/imitation/envs/examples/airl_envs/common.py b/src/imitation/envs/examples/airl_envs/common.py
deleted file mode 100644
index 9f7f051b1..000000000
--- a/src/imitation/envs/examples/airl_envs/common.py
+++ /dev/null
@@ -1,7 +0,0 @@
-import os
-
-ENV_ASSET_DIR = os.path.join(os.path.dirname(__file__), "assets")
-
-
-def get_asset_xml(xml_name):
-    return os.path.join(ENV_ASSET_DIR, xml_name)
diff --git a/src/imitation/envs/examples/airl_envs/dynamic_mjc/__init__.py b/src/imitation/envs/examples/airl_envs/dynamic_mjc/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/src/imitation/envs/examples/airl_envs/dynamic_mjc/mjc_models.py b/src/imitation/envs/examples/airl_envs/dynamic_mjc/mjc_models.py
deleted file mode 100644
index 781f7d6ee..000000000
--- a/src/imitation/envs/examples/airl_envs/dynamic_mjc/mjc_models.py
+++ /dev/null
@@ -1,1527 +0,0 @@
-import numpy as np
-
-from imitation.envs.examples.airl_envs.dynamic_mjc.model_builder import MJCModel
-
-
-def block_push(object_pos=(0, 0, 0), goal_pos=(0, 0, 0)):
-    mjcmodel = MJCModel("block_push")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="radian", coordinate="local")
-    mjcmodel.root.option(
-        timestep="0.01",
-        gravity="0 0 0",
-        iterations="20",
-        integrator="Euler",
-    )
-    default = mjcmodel.root.default()
-    default.joint(armature="0.04", damping=1, limited="true")
-    default.geom(
-        friction=".8 .1 .1",
-        density="300",
-        margin="0.002",
-        condim="1",
-        contype="1",
-        conaffinity="1",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-
-    palm = worldbody.body(name="palm", pos=[0, 0, 0])
-    palm.geom(
-        name="palm_geom",
-        type="capsule",
-        fromto=[0, -0.1, 0, 0, 0.1, 0],
-        size=0.12,
-    )
-    proximal1 = palm.body(name="proximal_1", pos=[0, 0, 0])
-    proximal1.joint(
-        name="proximal_j_1",
-        type="hinge",
-        pos=[0, 0, 0],
-        axis=[0, 1, 0],
-        range=[-2.5, 2.3],
-    )
-    proximal1.geom(
-        type="capsule",
-        fromto=[0, 0, 0, 0.4, 0, 0],
-        size=0.06,
-        contype=1,
-        conaffinity=1,
-    )
-    distal1 = proximal1.body(name="distal_1", pos=[0.4, 0, 0])
-    distal1.joint(
-        name="distal_j_1",
-        type="hinge",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-2.3213 2.3",
-        damping="1.0",
-    )
-    distal1.geom(
-        type="capsule",
-        fromto="0 0 0 0.4 0 0",
-        size="0.06",
-        contype="1",
-        conaffinity="1",
-    )
-    distal2 = distal1.body(name="distal_2", pos=[0.4, 0, 0])
-    distal2.joint(
-        name="distal_j_2",
-        type="hinge",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-2.3213 2.3",
-        damping="1.0",
-    )
-    distal2.geom(
-        type="capsule",
-        fromto="0 0 0 0.4 0 0",
-        size="0.06",
-        contype="1",
-        conaffinity="1",
-    )
-    distal4 = distal2.body(name="distal_4", pos=[0.4, 0, 0])
-    distal4.site(name="tip arml", pos="0.1 0 -0.2", size="0.01")
-    distal4.site(name="tip armr", pos="0.1 0 0.2", size="0.01")
-    distal4.joint(
-        name="distal_j_3",
-        type="hinge",
-        pos="0 0 0",
-        axis="1 0 0",
-        range="-3.3213 3.3",
-        damping="0.5",
-    )
-    distal4.geom(
-        type="capsule",
-        fromto="0 0 -0.2 0 0 0.2",
-        size="0.04",
-        contype="1",
-        conaffinity="1",
-    )
-    distal4.geom(
-        type="capsule",
-        fromto="0 0 -0.2 0.2 0 -0.2",
-        size="0.04",
-        contype="1",
-        conaffinity="1",
-    )
-    distal4.geom(
-        type="capsule",
-        fromto="0 0 0.2 0.2 0 0.2",
-        size="0.04",
-        contype="1",
-        conaffinity="1",
-    )
-
-    obj = worldbody.body(name="obj", pos=object_pos)
-    obj.geom(
-        rgba="1. 1. 1. 1",
-        type="box",
-        size="0.05 0.05 0.05",
-        density="0.00001",
-        contype="1",
-        conaffinity="1",
-    )
-    obj.joint(
-        name="obj_slidez",
-        type="slide",
-        pos="0.025 0.025 0.025",
-        axis="0 0 1",
-        range="-10.3213 10.3",
-        damping="0.5",
-    )
-    obj.joint(
-        name="obj_slidex",
-        type="slide",
-        pos="0.025 0.025 0.025",
-        axis="1 0 0",
-        range="-10.3213 10.3",
-        damping="0.5",
-    )
-    distal10 = obj.body(name="distal_10", pos=[0, 0, 0])
-    distal10.site(name="obj_pos", pos=[0.025, 0.025, 0.025], size=0.01)
-
-    goal = worldbody.body(name="goal", pos=goal_pos)
-    goal.geom(
-        rgba="1. 0. 0. 1",
-        type="box",
-        size="0.1 0.1 0.1",
-        density="0.00001",
-        contype="0",
-        conaffinity="0",
-    )
-    distal11 = goal.body(name="distal_11", pos=[0, 0, 0])
-    distal11.site(name="goal_pos", pos=[0.05, 0.05, 0.05], size=0.01)
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(joint="proximal_j_1", ctrlrange="-2 2", ctrllimited="true")
-    actuator.motor(joint="distal_j_1", ctrlrange="-2 2", ctrllimited="true")
-    actuator.motor(joint="distal_j_2", ctrlrange="-2 2", ctrllimited="true")
-    actuator.motor(joint="distal_j_3", ctrlrange="-2 2", ctrllimited="true")
-
-    return mjcmodel
-
-
-EAST = 0
-WEST = 1
-NORTH = 2
-SOUTH = 3
-
-
-def twod_corridor(direction=EAST, length=1.2):
-    mjcmodel = MJCModel("twod_corridor")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="radian", coordinate="local")
-    mjcmodel.root.option(
-        timestep="0.01",
-        gravity="0 0 0",
-        iterations="20",
-        integrator="Euler",
-    )
-    default = mjcmodel.root.default()
-    default.joint(damping=1, limited="false")
-    default.geom(
-        friction=".5 .1 .1",
-        density="1000",
-        margin="0.002",
-        condim="1",
-        contype="2",
-        conaffinity="1",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-
-    particle = worldbody.body(name="particle", pos=[0, 0, 0])
-    particle.geom(
-        name="particle_geom",
-        type="sphere",
-        size="0.03",
-        rgba="0.0 0.0 1.0 1",
-        contype=1,
-    )
-    particle.site(name="particle_site", pos=[0, 0, 0], size=0.01)
-    particle.joint(name="ball_x", type="slide", pos=[0, 0, 0], axis=[1, 0, 0])
-    particle.joint(name="ball_y", type="slide", pos=[0, 0, 0], axis=[0, 1, 0])
-
-    pos = np.array([0.0, 0, 0])
-    if direction == EAST or direction == WEST:
-        pos[0] = length - 0.1
-    else:
-        pos[1] = length - 0.1
-    if direction == WEST or direction == SOUTH:
-        pos = -pos
-
-    target = worldbody.body(name="target", pos=pos)
-    target.geom(
-        name="target_geom",
-        conaffinity=2,
-        type="sphere",
-        size=0.02,
-        rgba=[0, 0.9, 0.1, 1],
-    )
-
-    # arena
-    if direction == EAST:
-        L = -0.1
-        R = length
-        U = 0.1
-        D = -0.1
-    elif direction == WEST:
-        L = -length
-        R = 0.1
-        U = 0.1
-        D = -0.1
-    elif direction == SOUTH:
-        L = -0.1
-        R = 0.1
-        U = 0.1
-        D = -length
-    elif direction == NORTH:
-        L = -0.1
-        R = 0.1
-        U = length
-        D = -0.1
-
-    worldbody.geom(
-        conaffinity=1,
-        fromto=[L, D, 0.01, R, D, 0.01],
-        name="sideS",
-        rgba="0.9 0.4 0.6 1",
-        size=0.02,
-        type="capsule",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        fromto=[R, D, 0.01, R, U, 0.01],
-        name="sideE",
-        rgba="0.9 0.4 0.6 1",
-        size=".02",
-        type="capsule",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        fromto=[L, U, 0.01, R, U, 0.01],
-        name="sideN",
-        rgba="0.9 0.4 0.6 1",
-        size=".02",
-        type="capsule",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        fromto=[L, D, 0.01, L, U, 0.01],
-        name="sideW",
-        rgba="0.9 0.4 0.6 1",
-        size=".02",
-        type="capsule",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(joint="ball_x", ctrlrange=[-1.0, 1.0], ctrllimited=True)
-    actuator.motor(joint="ball_y", ctrlrange=[-1.0, 1.0], ctrllimited=True)
-
-    return mjcmodel
-
-
-LEFT = 0
-RIGHT = 1
-
-
-def point_mass_maze(direction=RIGHT, length=1.2, borders=True):
-    mjcmodel = MJCModel("twod_maze")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="radian", coordinate="local")
-    mjcmodel.root.option(
-        timestep="0.01",
-        gravity="0 0 0",
-        iterations="20",
-        integrator="Euler",
-    )
-    default = mjcmodel.root.default()
-    default.joint(damping=1, limited="false")
-    default.geom(
-        friction=".5 .1 .1",
-        density="1000",
-        margin="0.002",
-        condim="1",
-        contype="2",
-        conaffinity="1",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-
-    particle = worldbody.body(name="particle", pos=[length / 2, 0, 0])
-    particle.geom(
-        name="particle_geom",
-        type="sphere",
-        size="0.03",
-        rgba="0.0 0.0 1.0 1",
-        contype=1,
-    )
-    particle.site(name="particle_site", pos=[0, 0, 0], size=0.01)
-    particle.joint(name="ball_x", type="slide", pos=[0, 0, 0], axis=[1, 0, 0])
-    particle.joint(name="ball_y", type="slide", pos=[0, 0, 0], axis=[0, 1, 0])
-
-    target = worldbody.body(name="target", pos=[length / 2, length - 0.1, 0])
-    target.geom(
-        name="target_geom",
-        conaffinity=2,
-        type="sphere",
-        size=0.02,
-        rgba=[0, 0.9, 0.1, 1],
-    )
-
-    L = -0.1
-    R = length
-    U = length
-    D = -0.1
-
-    if borders:
-        worldbody.geom(
-            conaffinity=1,
-            fromto=[L, D, 0.01, R, D, 0.01],
-            name="sideS",
-            rgba="0.9 0.4 0.6 1",
-            size=".02",
-            type="capsule",
-        )
-        worldbody.geom(
-            conaffinity=1,
-            fromto=[R, D, 0.01, R, U, 0.01],
-            name="sideE",
-            rgba="0.9 0.4 0.6 1",
-            size=".02",
-            type="capsule",
-        )
-        worldbody.geom(
-            conaffinity=1,
-            fromto=[L, U, 0.01, R, U, 0.01],
-            name="sideN",
-            rgba="0.9 0.4 0.6 1",
-            size=".02",
-            type="capsule",
-        )
-        worldbody.geom(
-            conaffinity=1,
-            fromto=[L, D, 0.01, L, U, 0.01],
-            name="sideW",
-            rgba="0.9 0.4 0.6 1",
-            size=".02",
-            type="capsule",
-        )
-
-    # arena
-    if direction == LEFT:
-        BL = -0.1
-        BR = length * 2 / 3
-        BH = length / 2
-    else:
-        BL = length * 1 / 3
-        BR = length
-        BH = length / 2
-
-    worldbody.geom(
-        conaffinity=1,
-        fromto=[BL, BH, 0.01, BR, BH, 0.01],
-        name="barrier",
-        rgba="0.9 0.4 0.6 1",
-        size=".02",
-        type="capsule",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(joint="ball_x", ctrlrange=[-1.0, 1.0], ctrllimited=True)
-    actuator.motor(joint="ball_y", ctrlrange=[-1.0, 1.0], ctrllimited=True)
-
-    return mjcmodel
-
-
-def ant_maze(direction=RIGHT, length=6.0):
-    mjcmodel = MJCModel("ant_maze")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
-    mjcmodel.root.option(
-        timestep="0.01",
-        gravity="0 0 -9.8",
-        iterations="20",
-        integrator="Euler",
-    )
-
-    assets = mjcmodel.root.asset()
-    assets.texture(
-        builtin="gradient",
-        height="100",
-        rgb1="1 1 1",
-        rgb2="0 0 0",
-        type="skybox",
-        width="100",
-    )
-    assets.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",
-    )
-    assets.texture(
-        builtin="checker",
-        height="100",
-        name="texplane",
-        rgb1="0 0 0",
-        rgb2="0.8 0.8 0.8",
-        type="2d",
-        width="100",
-    )
-    assets.material(
-        name="MatPlane",
-        reflectance="0.5",
-        shininess="1",
-        specular="1",
-        texrepeat="60 60",
-        texture="texplane",
-    )
-    assets.material(name="geom", texture="texgeom", texuniform="true")
-
-    default = mjcmodel.root.default()
-    default.joint(armature="1", damping=1, limited="true")
-    default.geom(
-        friction="1 0.5 0.5",
-        density="5.0",
-        margin="0.01",
-        condim="3",
-        conaffinity="0",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-
-    ant = worldbody.body(name="ant", pos=[length / 2, 1.0, 0.05])
-    ant.geom(name="torso_geom", pos=[0, 0, 0], size="0.25", type="sphere")
-    ant.joint(
-        armature="0",
-        damping="0",
-        limited="false",
-        margin="0.01",
-        name="root",
-        pos=[0, 0, 0],
-        type="free",
-    )
-
-    front_left_leg = ant.body(name="front_left_leg", pos=[0, 0, 0])
-    front_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="aux_1_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_1 = front_left_leg.body(name="aux_1", pos=[0.2, 0.2, 0])
-    aux_1.joint(
-        axis=[0, 0, 1],
-        name="hip_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="left_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_1 = aux_1.body(pos=[0.2, 0.2, 0])
-    ankle_1.joint(
-        axis=[-1, 1, 0],
-        name="ankle_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, 0.4, 0.0],
-        name="left_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    front_right_leg = ant.body(name="front_right_leg", pos=[0, 0, 0])
-    front_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="aux_2_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_2 = front_right_leg.body(name="aux_2", pos=[-0.2, 0.2, 0])
-    aux_2.joint(
-        axis=[0, 0, 1],
-        name="hip_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="right_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_2 = aux_2.body(pos=[-0.2, 0.2, 0])
-    ankle_2.joint(
-        axis=[1, 1, 0],
-        name="ankle_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, 0.4, 0.0],
-        name="right_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    back_left_leg = ant.body(name="back_left_leg", pos=[0, 0, 0])
-    back_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="aux_3_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_3 = back_left_leg.body(name="aux_3", pos=[-0.2, -0.2, 0])
-    aux_3.joint(
-        axis=[0, 0, 1],
-        name="hip_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_3.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="backleft_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_3 = aux_3.body(pos=[-0.2, -0.2, 0])
-    ankle_3.joint(
-        axis=[-1, 1, 0],
-        name="ankle_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_3.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, -0.4, 0.0],
-        name="backleft_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    back_right_leg = ant.body(name="back_right_leg", pos=[0, 0, 0])
-    back_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="aux_4_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_4 = back_right_leg.body(name="aux_4", pos=[0.2, -0.2, 0])
-    aux_4.joint(
-        axis=[0, 0, 1],
-        name="hip_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_4.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="backright_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_4 = aux_4.body(pos=[0.2, -0.2, 0])
-    ankle_4.joint(
-        axis=[1, 1, 0],
-        name="ankle_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_4.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, -0.4, 0.0],
-        name="backright_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    target = worldbody.body(name="target", pos=[length / 2, length - 0.2, -0.5])
-    target.geom(
-        name="target_geom",
-        conaffinity=2,
-        type="sphere",
-        size=0.2,
-        rgba=[0, 0.9, 0.1, 1],
-    )
-
-    L = length / 2
-    h = 0.75
-    w = 0.05
-
-    worldbody.geom(
-        conaffinity=1,
-        name="sideS",
-        rgba="0.9 0.4 0.6 1",
-        size=[L, w, h],
-        pos=[length / 2, 0, 0],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        name="sideE",
-        rgba="0.9 0.4 0.6 1",
-        size=[w, L, h],
-        pos=[length, length / 2, 0],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        name="sideN",
-        rgba="0.9 0.4 0.6 1",
-        size=[L, w, h],
-        pos=[length / 2, length, 0],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        name="sideW",
-        rgba="0.9 0.4 0.6 1",
-        size=[w, L, h],
-        pos=[0, length / 2, 0],
-        type="box",
-    )
-
-    # arena
-    if direction == LEFT:
-        bx, by, bz = (length / 3, length / 2, 0)
-    else:
-        bx, by, bz = (length * 2 / 3, length / 2, 0)
-
-    worldbody.geom(
-        conaffinity=1,
-        name="barrier",
-        rgba="0.9 0.4 0.6 1",
-        size=[L * 2 / 3, w, h],
-        pos=[bx, by, bz],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity="1",
-        condim="3",
-        material="MatPlane",
-        name="floor",
-        pos=[length / 2, length / 2, -h + w],
-        rgba="0.8 0.9 0.8 1",
-        size="40 40 40",
-        type="plane",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_4", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_4", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_1", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_1", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_2", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_2", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_3", gear="50")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_3", gear="50")
-
-    return mjcmodel
-
-
-def ant_maze_corridor(direction=RIGHT, height=6.0, width=10.0):
-    mjcmodel = MJCModel("ant_maze_corridor")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
-    mjcmodel.root.option(
-        timestep="0.01",
-        gravity="0 0 -9.8",
-        iterations="20",
-        integrator="Euler",
-    )
-
-    assets = mjcmodel.root.asset()
-    assets.texture(
-        builtin="gradient",
-        height="100",
-        rgb1="1 1 1",
-        rgb2="0 0 0",
-        type="skybox",
-        width="100",
-    )
-    assets.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",
-    )
-    assets.texture(
-        builtin="checker",
-        height="100",
-        name="texplane",
-        rgb1="0 0 0",
-        rgb2="0.8 0.8 0.8",
-        type="2d",
-        width="100",
-    )
-    assets.material(
-        name="MatPlane",
-        reflectance="0.5",
-        shininess="1",
-        specular="1",
-        texrepeat="60 60",
-        texture="texplane",
-    )
-    assets.material(name="geom", texture="texgeom", texuniform="true")
-
-    default = mjcmodel.root.default()
-    default.joint(armature="1", damping=1, limited="true")
-    default.geom(
-        friction="1 0.5 0.5",
-        density="5.0",
-        margin="0.01",
-        condim="3",
-        conaffinity="0",
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-
-    ant = worldbody.body(name="ant", pos=[height / 2, 1.0, 0.05])
-    ant.geom(name="torso_geom", pos=[0, 0, 0], size="0.25", type="sphere")
-    ant.joint(
-        armature="0",
-        damping="0",
-        limited="false",
-        margin="0.01",
-        name="root",
-        pos=[0, 0, 0],
-        type="free",
-    )
-
-    front_left_leg = ant.body(name="front_left_leg", pos=[0, 0, 0])
-    front_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="aux_1_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_1 = front_left_leg.body(name="aux_1", pos=[0.2, 0.2, 0])
-    aux_1.joint(
-        axis=[0, 0, 1],
-        name="hip_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, 0.2, 0.0],
-        name="left_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_1 = aux_1.body(pos=[0.2, 0.2, 0])
-    ankle_1.joint(
-        axis=[-1, 1, 0],
-        name="ankle_1",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_1.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, 0.4, 0.0],
-        name="left_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    front_right_leg = ant.body(name="front_right_leg", pos=[0, 0, 0])
-    front_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="aux_2_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_2 = front_right_leg.body(name="aux_2", pos=[-0.2, 0.2, 0])
-    aux_2.joint(
-        axis=[0, 0, 1],
-        name="hip_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, 0.2, 0.0],
-        name="right_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_2 = aux_2.body(pos=[-0.2, 0.2, 0])
-    ankle_2.joint(
-        axis=[1, 1, 0],
-        name="ankle_2",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_2.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, 0.4, 0.0],
-        name="right_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    back_left_leg = ant.body(name="back_left_leg", pos=[0, 0, 0])
-    back_left_leg.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="aux_3_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_3 = back_left_leg.body(name="aux_3", pos=[-0.2, -0.2, 0])
-    aux_3.joint(
-        axis=[0, 0, 1],
-        name="hip_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_3.geom(
-        fromto=[0.0, 0.0, 0.0, -0.2, -0.2, 0.0],
-        name="backleft_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_3 = aux_3.body(pos=[-0.2, -0.2, 0])
-    ankle_3.joint(
-        axis=[-1, 1, 0],
-        name="ankle_3",
-        pos=[0.0, 0.0, 0.0],
-        range=[-70, -30],
-        type="hinge",
-    )
-    ankle_3.geom(
-        fromto=[0.0, 0.0, 0.0, -0.4, -0.4, 0.0],
-        name="backleft_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    back_right_leg = ant.body(name="back_right_leg", pos=[0, 0, 0])
-    back_right_leg.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="aux_4_geom",
-        size="0.08",
-        type="capsule",
-    )
-    aux_4 = back_right_leg.body(name="aux_4", pos=[0.2, -0.2, 0])
-    aux_4.joint(
-        axis=[0, 0, 1],
-        name="hip_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[-30, 30],
-        type="hinge",
-    )
-    aux_4.geom(
-        fromto=[0.0, 0.0, 0.0, 0.2, -0.2, 0.0],
-        name="backright_leg_geom",
-        size="0.08",
-        type="capsule",
-    )
-    ankle_4 = aux_4.body(pos=[0.2, -0.2, 0])
-    ankle_4.joint(
-        axis=[1, 1, 0],
-        name="ankle_4",
-        pos=[0.0, 0.0, 0.0],
-        range=[30, 70],
-        type="hinge",
-    )
-    ankle_4.geom(
-        fromto=[0.0, 0.0, 0.0, 0.4, -0.4, 0.0],
-        name="backright_ankle_geom",
-        size="0.08",
-        type="capsule",
-    )
-
-    target = worldbody.body(name="target", pos=[height / 2, width - 1.0, -0.5])
-    target.geom(
-        name="target_geom",
-        conaffinity=2,
-        type="sphere",
-        size=0.2,
-        rgba=[0, 0.9, 0.1, 1],
-    )
-
-    L = height / 2
-    h = 0.75
-    w = 0.05
-
-    worldbody.geom(
-        conaffinity=1,
-        name="sideS",
-        rgba="0.9 0.4 0.6 1",
-        size=[L, w, h],
-        pos=[height / 2, 0, 0],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        name="sideE",
-        rgba="0.9 0.4 0.6 1",
-        size=[w, width / 2, h],
-        pos=[height, width / 2, 0],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        name="sideN",
-        rgba="0.9 0.4 0.6 1",
-        size=[L, w, h],
-        pos=[height / 2, width, 0],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity=1,
-        name="sideW",
-        rgba="0.9 0.4 0.6 1",
-        size=[w, width / 2, h],
-        pos=[0, width / 2, 0],
-        type="box",
-    )
-
-    # arena
-    wall_ratio = 0.55  # 2.0/3
-    if direction == LEFT:
-        bx, by, bz = (height * (wall_ratio / 2), width / 2, 0)
-    else:
-        bx, by, bz = (height * (1 - wall_ratio / 2), width / 2, 0)
-
-    worldbody.geom(
-        conaffinity=1,
-        name="barrier",
-        rgba="0.9 0.4 0.6 1",
-        size=[L * (wall_ratio), w, h],
-        pos=[bx, by, bz],
-        type="box",
-    )
-    worldbody.geom(
-        conaffinity="1",
-        condim="3",
-        material="MatPlane",
-        name="floor",
-        pos=[height / 2, height / 2, -h + w],
-        rgba="0.8 0.9 0.8 1",
-        size="40 40 40",
-        type="plane",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_4", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_4", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_1", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_1", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_2", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_2", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="hip_3", gear="30")
-    actuator.motor(ctrllimited="true", ctrlrange="-1.0 1.0", joint="ankle_3", gear="30")
-
-    return mjcmodel
-
-
-def pusher(goal_pos=np.array([0.45, -0.05, -0.323])):
-    mjcmodel = MJCModel("pusher")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="radian", coordinate="local")
-    mjcmodel.root.option(
-        timestep="0.01",
-        gravity="0 0 0",
-        iterations="20",
-        integrator="Euler",
-    )
-    default = mjcmodel.root.default()
-    default.joint(armature=0.04, damping=1, limited=False)
-    default.geom(
-        friction=[0.8, 0.1, 0.1],
-        density=300,
-        margin=0.002,
-        condim=1,
-        contype=0,
-        conaffinity=0,
-    )
-
-    worldbody = mjcmodel.root.worldbody()
-    worldbody.light(diffuse=[0.5, 0.5, 0.5], pos=[0, 0, 3], dir=[0, 0, -1])
-    worldbody.geom(
-        name="table",
-        type="plane",
-        pos=[0, 0.5, -0.325],
-        size=[1, 1, 0.1],
-        contype=1,
-        conaffinity=1,
-    )
-
-    r_shoulder_pan_link = worldbody.body(name="r_shoulder_pan_link", pos=[0, -0.6, 0])
-    r_shoulder_pan_link.geom(
-        name="e1",
-        type="sphere",
-        rgba=[0.6, 0.6, 0.6, 1],
-        pos=[-0.06, 0.05, 0.2],
-        size=0.05,
-    )
-    r_shoulder_pan_link.geom(
-        name="e2",
-        type="sphere",
-        rgba=[0.6, 0.6, 0.6, 1],
-        pos=[0.06, 0.05, 0.2],
-        size=0.05,
-    )
-    r_shoulder_pan_link.geom(
-        name="e1p",
-        type="sphere",
-        rgba=[0.1, 0.1, 0.1, 1],
-        pos=[-0.06, 0.09, 0.2],
-        size=0.03,
-    )
-    r_shoulder_pan_link.geom(
-        name="e2p",
-        type="sphere",
-        rgba=[0.1, 0.1, 0.1, 1],
-        pos=[0.06, 0.09, 0.2],
-        size=0.03,
-    )
-    r_shoulder_pan_link.geom(
-        name="sp",
-        type="capsule",
-        fromto=[0, 0, -0.4, 0, 0, 0.2],
-        size=0.1,
-    )
-
-    r_shoulder_pan_link.joint(
-        name="r_shoulder_pan_joint",
-        type="hinge",
-        pos=[0, 0, 0],
-        axis=[0, 0, 1],
-        range=[-2.2854, 1.714602],
-        damping=1.0,
-    )
-
-    r_shoulder_lift_link = r_shoulder_pan_link.body(
-        name="r_shoulder_lift_link",
-        pos=[0.1, 0, 0],
-    )
-    r_shoulder_lift_link.geom(
-        name="s1",
-        type="capsule",
-        fromto="0 -0.1 0 0 0.1 0",
-        size="0.1",
-    )
-    r_shoulder_lift_link.joint(
-        name="r_shoulder_lift_joint",
-        type="hinge",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-0.5236 1.3963",
-        damping="1.0",
-    )
-
-    r_upper_arm_roll_link = r_shoulder_lift_link.body(
-        name="r_upper_arm_roll_link",
-        pos=[0, 0, 0],
-    )
-    r_upper_arm_roll_link.geom(
-        name="uar",
-        type="capsule",
-        fromto="-0.1 0 0 0.1 0 0",
-        size="0.02",
-    )
-    r_upper_arm_roll_link.joint(
-        name="r_upper_arm_roll_joint",
-        type="hinge",
-        pos="0 0 0",
-        axis="1 0 0",
-        range="-1.5 1.7",
-        damping="0.1",
-    )
-
-    r_upper_arm_link = r_upper_arm_roll_link.body(
-        name="r_upper_arm_link",
-        pos=[0, 0, 0],
-    )
-    r_upper_arm_link.geom(
-        name="ua",
-        type="capsule",
-        fromto="0 0 0 0.4 0 0",
-        size="0.06",
-    )
-
-    r_elbow_flex_link = r_upper_arm_link.body(name="r_elbow_flex_link", pos=[0.4, 0, 0])
-    r_elbow_flex_link.geom(
-        name="ef",
-        type="capsule",
-        fromto="0 -0.02 0 0.0 0.02 0",
-        size="0.06",
-    )
-    r_elbow_flex_link.joint(
-        name="r_elbow_flex_joint",
-        type="hinge",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-2.3213 0",
-        damping="0.1",
-    )
-
-    r_forearm_roll_link = r_elbow_flex_link.body(
-        name="r_forearm_roll_link",
-        pos=[0, 0, 0],
-    )
-    r_forearm_roll_link.geom(
-        name="fr",
-        type="capsule",
-        fromto="-0.1 0 0 0.1 0 0",
-        size="0.02",
-    )
-    r_forearm_roll_link.joint(
-        name="r_forearm_roll_joint",
-        type="hinge",
-        limited="true",
-        pos="0 0 0",
-        axis="1 0 0",
-        damping=".1",
-        range="-1.5 1.5",
-    )
-
-    r_forearm_link = r_forearm_roll_link.body(name="r_forearm_link", pos=[0, 0, 0])
-    r_forearm_link.geom(
-        name="fa",
-        type="capsule",
-        fromto="0 0 0 0.291 0 0",
-        size="0.05",
-    )
-
-    r_wrist_flex_link = r_forearm_link.body(name="r_wrist_flex_link", pos=[0.321, 0, 0])
-    r_wrist_flex_link.geom(
-        name="wf",
-        type="capsule",
-        fromto="0 -0.02 0 0 0.02 0",
-        size="0.01",
-    )
-    r_wrist_flex_link.joint(
-        name="r_wrist_flex_joint",
-        type="hinge",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-1.094 0",
-        damping=".1",
-    )
-
-    r_wrist_roll_link = r_wrist_flex_link.body(name="r_wrist_roll_link", pos=[0, 0, 0])
-    r_wrist_roll_link.joint(
-        name="r_wrist_roll_joint",
-        type="hinge",
-        pos="0 0 0",
-        limited="true",
-        axis="1 0 0",
-        damping="0.1",
-        range="-1.5 1.5",
-    )
-    r_wrist_roll_link.geom(
-        type="capsule",
-        fromto="0 -0.1 0. 0.0 +0.1 0",
-        size="0.02",
-        contype="1",
-        conaffinity="1",
-    )
-    r_wrist_roll_link.geom(
-        type="capsule",
-        fromto="0 -0.1 0. 0.1 -0.1 0",
-        size="0.02",
-        contype="1",
-        conaffinity="1",
-    )
-    r_wrist_roll_link.geom(
-        type="capsule",
-        fromto="0 +0.1 0. 0.1 +0.1 0",
-        size="0.02",
-        contype="1",
-        conaffinity="1",
-    )
-
-    tips_arm = r_wrist_roll_link.body(name="tips_arm", pos=[0, 0, 0])
-    tips_arm.geom(name="tip_arml", type="sphere", pos="0.1 -0.1 0.", size="0.01")
-    tips_arm.geom(name="tip_armr", type="sphere", pos="0.1 0.1 0.", size="0.01")
-
-    object_ = worldbody.body(name="object", pos=[0.0, 0.0, -0.275])
-    object_.geom(
-        rgba="1 1 1 1",
-        type="cylinder",
-        size="0.05 0.05 0.05",
-        density="0.00001",
-        conaffinity="0",
-        contype=1,
-    )
-    object_.joint(
-        name="obj_slidey",
-        type="slide",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-10.3213 10.3",
-        damping="0.5",
-    )
-    object_.joint(
-        name="obj_slidex",
-        type="slide",
-        pos="0 0 0",
-        axis="1 0 0",
-        range="-10.3213 10.3",
-        damping="0.5",
-    )
-
-    goal = worldbody.body(name="goal", pos=goal_pos)
-    goal.geom(
-        rgba="1 0 0 1",
-        type="cylinder",
-        size="0.08 0.001 0.1",
-        density="0.00001",
-        contype="0",
-        conaffinity="0",
-    )
-    goal.joint(
-        name="goal_slidey",
-        type="slide",
-        pos="0 0 0",
-        axis="0 1 0",
-        range="-10.3213 10.3",
-        damping="0.5",
-    )
-    goal.joint(
-        name="goal_slidex",
-        type="slide",
-        pos="0 0 0",
-        axis="1 0 0",
-        range="-10.3213 10.3",
-        damping="0.5",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(
-        joint="r_shoulder_pan_joint",
-        ctrlrange=[-2.0, 2.0],
-        ctrllimited=True,
-    )
-    actuator.motor(
-        joint="r_shoulder_lift_joint",
-        ctrlrange=[-2.0, 2.0],
-        ctrllimited=True,
-    )
-    actuator.motor(
-        joint="r_upper_arm_roll_joint",
-        ctrlrange=[-2.0, 2.0],
-        ctrllimited=True,
-    )
-    actuator.motor(joint="r_elbow_flex_joint", ctrlrange=[-2.0, 2.0], ctrllimited=True)
-    actuator.motor(
-        joint="r_forearm_roll_joint",
-        ctrlrange=[-2.0, 2.0],
-        ctrllimited=True,
-    )
-    actuator.motor(joint="r_wrist_flex_joint", ctrlrange=[-2.0, 2.0], ctrllimited=True)
-    actuator.motor(joint="r_wrist_roll_joint", ctrlrange=[-2.0, 2.0], ctrllimited=True)
-
-    return mjcmodel
-
-
-def swimmer():
-    mjcmodel = MJCModel("swimmer")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
-    mjcmodel.root.option(
-        timestep=0.01,
-        viscosity=0.1,
-        density=4000,
-        integrator="RK4",
-        collision="predefined",
-    )
-    default = mjcmodel.root.default()
-    default.joint(armature=0.1)
-    default.geom(
-        rgba=[0.8, 0.6, 0.1, 1],
-        condim=1,
-        contype=1,
-        conaffinity=1,
-        material="geom",
-    )
-
-    asset = mjcmodel.root.asset()
-    asset.texture(
-        builtin="gradient",
-        height=100,
-        rgb1=[1, 1, 1],
-        rgb2=[0, 0, 0],
-        type="skybox",
-        width=100,
-    )
-    asset.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,
-    )
-    asset.texture(
-        builtin="checker",
-        height=100,
-        name="texplane",
-        rgb1=[0, 0, 0],
-        rgb2=[0.8, 0.8, 0.8],
-        type="2d",
-        width=100,
-    )
-    asset.material(
-        name="MatPlane",
-        reflectance=0.5,
-        shininess=1,
-        specular=1,
-        texrepeat=[30, 30],
-        texture="texplane",
-    )
-    asset.material(name="geom", texture="texgeom", texuniform=True)
-
-    worldbody = mjcmodel.root.worldbody()
-    worldbody.light(
-        cutoff=100,
-        diffuse=[1, 1, 1],
-        dir=[0, 0, -1.3],
-        directional=True,
-        exponent=1,
-        pos=[0, 0, 1.3],
-        specular=[0.1, 0.1, 0.1],
-    )
-    worldbody.geom(
-        conaffinity=1,
-        condim=3,
-        material="MatPlane",
-        name="floor",
-        pos=[0, 0, -0.1],
-        rgba=[0.8, 0.9, 0.9, 1],
-        size=[40, 40, 0.1],
-        type="plane",
-    )
-    torso = worldbody.body(name="torso", pos=[0, 0, 0])
-    torso.geom(density=1000, fromto=[1.5, 0, 0, 0.5, 0, 0], size=0.1, type="capsule")
-    torso.joint(axis=[1, 0, 0], name="slider1", pos=[0, 0, 0], type="slide")
-    torso.joint(axis=[0, 1, 0], name="slider2", pos=[0, 0, 0], type="slide")
-    torso.joint(axis=[0, 0, 1], name="rot", pos=[0, 0, 0], type="hinge")
-    mid = torso.body(name="mid", pos=[0.5, 0, 0])
-    mid.geom(density=1000, fromto=[0, 0, 0, -1, 0, 0], size=0.1, type="capsule")
-    mid.joint(
-        axis=[0, 0, 1],
-        limited=True,
-        name="rot2",
-        pos=[0, 0, 0],
-        range=[-100, 100],
-        type="hinge",
-    )
-    back = mid.body(name="back", pos=[-1, 0, 0])
-    back.geom(density=1000, fromto=[0, 0, 0, -1, 0, 0], size=0.1, type="capsule")
-    back.joint(
-        axis=[0, 0, 1],
-        limited=True,
-        name="rot3",
-        pos=[0, 0, 0],
-        range=[-100, 100],
-        type="hinge",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(ctrllimited=True, ctrlrange=[-1, 1], gear=150, joint="rot2")
-    actuator.motor(ctrllimited=True, ctrlrange=[-1, 1], gear=150, joint="rot3")
-    return mjcmodel
-
-
-def swimmer_rllab():
-    mjcmodel = MJCModel("swimmer")
-    mjcmodel.root.compiler(inertiafromgeom="true", angle="degree", coordinate="local")
-    mjcmodel.root.option(
-        timestep=0.01,
-        viscosity=0.1,
-        density=4000,
-        integrator="Euler",
-        iterations=1000,
-        collision="predefined",
-    )
-
-    custom = mjcmodel.root.custom()
-    custom.numeric(name="frame_skip", data=50)
-
-    default = mjcmodel.root.default()
-    # default.joint(armature=0.1)
-    default.geom(
-        rgba=[0.8, 0.6, 0.1, 1],
-        condim=1,
-        contype=1,
-        conaffinity=1,
-        material="geom",
-    )
-
-    asset = mjcmodel.root.asset()
-    asset.texture(
-        builtin="gradient",
-        height=100,
-        rgb1=[1, 1, 1],
-        rgb2=[0, 0, 0],
-        type="skybox",
-        width=100,
-    )
-    asset.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,
-    )
-    asset.texture(
-        builtin="checker",
-        height=100,
-        name="texplane",
-        rgb1=[0, 0, 0],
-        rgb2=[0.8, 0.8, 0.8],
-        type="2d",
-        width=100,
-    )
-    asset.material(
-        name="MatPlane",
-        reflectance=0.5,
-        shininess=1,
-        specular=1,
-        texrepeat=[30, 30],
-        texture="texplane",
-    )
-    asset.material(name="geom", texture="texgeom", texuniform=True)
-
-    worldbody = mjcmodel.root.worldbody()
-    worldbody.light(
-        cutoff=100,
-        diffuse=[1, 1, 1],
-        dir=[0, 0, -1.3],
-        directional=True,
-        exponent=1,
-        pos=[0, 0, 1.3],
-        specular=[0.1, 0.1, 0.1],
-    )
-    worldbody.geom(
-        conaffinity=1,
-        condim=3,
-        material="MatPlane",
-        name="floor",
-        pos=[0, 0, -0.1],
-        rgba=[0.8, 0.9, 0.9, 1],
-        size=[40, 40, 0.1],
-        type="plane",
-    )
-    torso = worldbody.body(name="torso", pos=[0, 0, 0])
-    torso.geom(density=1000, fromto=[1.5, 0, 0, 0.5, 0, 0], size=0.1, type="capsule")
-    torso.joint(axis=[1, 0, 0], name="slider1", pos=[0, 0, 0], type="slide")
-    torso.joint(axis=[0, 1, 0], name="slider2", pos=[0, 0, 0], type="slide")
-    torso.joint(axis=[0, 0, 1], name="rot", pos=[0, 0, 0], type="hinge")
-    mid = torso.body(name="mid", pos=[0.5, 0, 0])
-    mid.geom(density=1000, fromto=[0, 0, 0, -1, 0, 0], size=0.1, type="capsule")
-    mid.joint(
-        axis=[0, 0, 1],
-        limited=True,
-        name="rot2",
-        pos=[0, 0, 0],
-        range=[-100, 100],
-        type="hinge",
-    )
-    back = mid.body(name="back", pos=[-1, 0, 0])
-    back.geom(density=1000, fromto=[0, 0, 0, -1, 0, 0], size=0.1, type="capsule")
-    back.joint(
-        axis=[0, 0, 1],
-        limited=True,
-        name="rot3",
-        pos=[0, 0, 0],
-        range=[-100, 100],
-        type="hinge",
-    )
-
-    actuator = mjcmodel.root.actuator()
-    actuator.motor(ctrllimited=True, ctrlrange=[-50, 50], joint="rot2")
-    actuator.motor(ctrllimited=True, ctrlrange=[-50, 50], joint="rot3")
-    return mjcmodel
diff --git a/src/imitation/envs/examples/airl_envs/dynamic_mjc/model_builder.py b/src/imitation/envs/examples/airl_envs/dynamic_mjc/model_builder.py
deleted file mode 100644
index 56114b945..000000000
--- a/src/imitation/envs/examples/airl_envs/dynamic_mjc/model_builder.py
+++ /dev/null
@@ -1,130 +0,0 @@
-"""model_builder.py A small library for programmatically building MuJoCo XML files."""
-import tempfile
-from contextlib import contextmanager
-
-import numpy as np
-
-
-def default_model(name):
-    """Get a model with basic settings such as gravity and RK4 integration."""
-    model = MJCModel(name)
-    root = model.root
-
-    # Setup
-    root.compiler(angle="radian", inertiafromgeom="true")
-    default = root.default()
-    default.joint(armature=1, damping=1, limited="true")
-    default.geom(contype=0, friction="1 0.1 0.1", rgba="0.7 0.7 0 1")
-    root.option(gravity="0 0 -9.81", integrator="RK4", timestep=0.01)
-    return model
-
-
-def pointmass_model(name):
-    """Get a model with basic settings such as gravity and Euler integration."""
-    model = MJCModel(name)
-    root = model.root
-
-    # Setup
-    root.compiler(angle="radian", inertiafromgeom="true", coordinate="local")
-    default = root.default()
-    default.joint(limited="false", damping=1)
-    default.geom(
-        contype=2,
-        conaffinity="1",
-        condim="1",
-        friction=".5 .1 .1",
-        density="1000",
-        margin="0.002",
-    )
-    root.option(timestep=0.01, gravity="0 0 0", iterations="20", integrator="Euler")
-    return model
-
-
-class MJCModel(object):
-    def __init__(self, name):
-        self.name = name
-        self.root = MJCTreeNode("mujoco").add_attr("model", name)
-
-    @contextmanager
-    def asfile(self):
-        """Usage:
-        model = MJCModel('reacher')
-        with model.asfile() as f:
-            print f.read()  # prints a dump of the model
-        """
-        with tempfile.NamedTemporaryFile(mode="w+", suffix=".xml", delete=True) as f:
-            self.root.write(f)
-            f.seek(0)
-            yield f
-
-    def find_attr(self, attr, value):
-        return self.root.find_attr(attr, value)
-
-    def __getstate__(self):
-        return {}
-
-    def __setstate__(self, state):
-        pass
-
-
-class MJCTreeNode(object):
-    def __init__(self, name):
-        self.name = name
-        self.attrs = {}
-        self.children = []
-
-    def add_attr(self, key, value):
-        if isinstance(value, str):
-            pass
-        elif isinstance(value, list) or isinstance(value, np.ndarray):
-            value = " ".join([str(val).lower() for val in value])
-        else:
-            value = str(value).lower()
-
-        self.attrs[key] = value
-        return self
-
-    def __getattr__(self, name):
-        def wrapper(**kwargs):
-            newnode = MJCTreeNode(name)
-            for (k, v) in kwargs.items():
-                newnode.add_attr(k, v)
-            self.children.append(newnode)
-            return newnode
-
-        return wrapper
-
-    def dfs(self):
-        yield self
-        if self.children:
-            for child in self.children:
-                for node in child.dfs():
-                    yield node
-
-    def find_attr(self, attr, value):
-        """Run DFS to find a matching attr."""
-        if attr in self.attrs and self.attrs[attr] == value:
-            return self
-        for child in self.children:
-            res = child.find_attr(attr, value)
-            if res is not None:
-                return res
-        return None
-
-    def write(self, ostream, tabs=0):
-        contents = " ".join(['%s="%s"' % (k, v) for (k, v) in self.attrs.items()])
-        if self.children:
-            ostream.write("\t" * tabs)
-            ostream.write("<%s %s>\n" % (self.name, contents))
-            for child in self.children:
-                child.write(ostream, tabs=tabs + 1)
-            ostream.write("\t" * tabs)
-            ostream.write("</%s>\n" % self.name)
-        else:
-            ostream.write("\t" * tabs)
-            ostream.write("<%s %s/>\n" % (self.name, contents))
-
-    def __str__(self):
-        s = "<" + self.name
-        s += " ".join(['%s="%s"' % (k, v) for (k, v) in self.attrs.items()])
-        return s + ">"
diff --git a/src/imitation/envs/examples/airl_envs/point_maze_env.py b/src/imitation/envs/examples/airl_envs/point_maze_env.py
deleted file mode 100644
index 8a3da1e18..000000000
--- a/src/imitation/envs/examples/airl_envs/point_maze_env.py
+++ /dev/null
@@ -1,75 +0,0 @@
-import numpy as np
-from gym import utils
-from gym.envs.mujoco import mujoco_env
-
-from imitation.envs.examples.airl_envs.dynamic_mjc.mjc_models import point_mass_maze
-
-
-class PointMazeEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(
-        self,
-        direction=1,
-        maze_length=0.6,
-        sparse_reward=False,
-        no_reward=False,
-        include_vel=False,
-        episode_length=100,
-    ):
-        utils.EzPickle.__init__(self)
-        self.sparse_reward = sparse_reward
-        self.no_reward = no_reward
-        self.include_vel = include_vel
-        self.max_episode_length = episode_length
-        self.direction = direction
-        self.length = maze_length
-
-        self.episode_length = 0
-
-        model = point_mass_maze(direction=self.direction, length=self.length)
-        with model.asfile() as f:
-            mujoco_env.MujocoEnv.__init__(self, f.name, 5)
-
-    def step(self, a):
-        vec_dist = self.get_body_com("particle") - self.get_body_com("target")
-
-        reward_dist = -np.linalg.norm(vec_dist)  # particle to target
-        reward_ctrl = -np.square(a).sum()
-        if self.no_reward:
-            reward = 0
-        elif self.sparse_reward:
-            if reward_dist <= 0.1:
-                reward = 1
-            else:
-                reward = 0
-        else:
-            reward = reward_dist + 0.001 * reward_ctrl
-
-        self.do_simulation(a, self.frame_skip)
-        ob = self._get_obs()
-        self.episode_length += 1
-        done = self.episode_length >= self.max_episode_length
-        return ob, reward, done, dict(reward_dist=reward_dist, reward_ctrl=reward_ctrl)
-
-    def viewer_setup(self):
-        self.viewer.cam.trackbodyid = -1
-        self.viewer.cam.distance = 4.0
-
-    def reset_model(self):
-        qpos = self.init_qpos
-        qvel = self.init_qvel + self.np_random.uniform(
-            size=self.model.nv,
-            low=-0.01,
-            high=0.01,
-        )
-        self.set_state(qpos, qvel)
-        self.episode_length = 0
-        return self._get_obs()
-
-    def _get_obs(self):
-        obs = [self.get_body_com("particle")]
-        if self.include_vel:
-            obs.append(self.sim.data.get_body_xvelp("particle"))
-        return np.concatenate(obs)
-
-    def plot_trajs(self, *args, **kwargs):
-        pass
diff --git a/src/imitation/envs/examples/airl_envs/pusher_env.py b/src/imitation/envs/examples/airl_envs/pusher_env.py
deleted file mode 100644
index dc5dc0381..000000000
--- a/src/imitation/envs/examples/airl_envs/pusher_env.py
+++ /dev/null
@@ -1,94 +0,0 @@
-import numpy as np
-from gym import utils
-from gym.envs.mujoco import mujoco_env
-
-from imitation.envs.examples.airl_envs.dynamic_mjc.mjc_models import pusher
-
-
-class PusherEnv(mujoco_env.MujocoEnv, utils.EzPickle):
-    def __init__(self, sparse_reward=False, no_reward=False, episode_length=200):
-        utils.EzPickle.__init__(self)
-        self.sparse_reward = sparse_reward
-        self.no_reward = no_reward
-        self.max_episode_length = episode_length
-        self.goal_pos = np.asarray([0.0, 0.0])
-
-        self.episode_length = 0
-
-        model = pusher(goal_pos=[self.goal_pos[0], self.goal_pos[1], -0.323])
-        with model.asfile() as f:
-            mujoco_env.MujocoEnv.__init__(self, f.name, 5)
-
-    def step(self, a):
-        vec_1 = self.get_body_com("object") - self.get_body_com("tips_arm")
-        vec_2 = self.get_body_com("object") - self.get_body_com("goal")
-
-        reward_near = -np.linalg.norm(vec_1)  # arm to object
-        reward_dist = -np.linalg.norm(vec_2[0:2])  # object to goal
-        reward_ctrl = -np.square(a).sum()
-        if self.no_reward:
-            reward = 0
-        elif self.sparse_reward:
-            reward = reward_dist + 0.1 * reward_ctrl
-        else:
-            reward = reward_dist + 0.1 * reward_ctrl + 0.5 * reward_near
-
-        self.do_simulation(a, self.frame_skip)
-        ob = self._get_obs()
-        self.episode_length += 1
-        done = self.episode_length >= self.max_episode_length
-        return (
-            ob,
-            reward,
-            done,
-            dict(
-                reward_dist=reward_dist,
-                reward_near=reward_near,
-                reward_ctrl=reward_ctrl,
-            ),
-        )
-
-    def viewer_setup(self):
-        self.viewer.cam.trackbodyid = -1
-        self.viewer.cam.distance = 4.0
-
-    def reset_model(self):
-        qpos = self.init_qpos
-        self.episode_length = 0
-
-        while True:
-            self.cylinder_pos = np.concatenate(
-                [
-                    self.np_random.uniform(low=-0.5, high=0, size=1),
-                    self.np_random.uniform(low=-0.5, high=0.5, size=1),
-                ],
-            )
-            cyl_dist = np.linalg.norm(self.cylinder_pos - self.goal_pos)
-            if cyl_dist > 0.2 and cyl_dist < 0.4:
-                break
-
-        qpos[-4:-2] = self.cylinder_pos
-        qpos[-2] = self.goal_pos[0]
-        qpos[-1] = self.goal_pos[1]
-        qvel = self.init_qvel + self.np_random.uniform(
-            low=-0.005,
-            high=0.005,
-            size=self.model.nv,
-        )
-        qvel[-4:] = 0
-        self.set_state(qpos, qvel)
-        return self._get_obs()
-
-    def _get_obs(self):
-        return np.concatenate(
-            [
-                self.sim.data.qpos.flat[:7],
-                self.sim.data.qvel.flat[:7],
-                self.get_body_com("tips_arm"),
-                self.get_body_com("object"),
-                self.get_body_com("goal"),
-            ],
-        )
-
-    def plot_trajs(self, *args, **kwargs):
-        pass
diff --git a/src/imitation/envs/examples/airl_envs/twod_maze.py b/src/imitation/envs/examples/airl_envs/twod_maze.py
deleted file mode 100644
index 5ef4b55fa..000000000
--- a/src/imitation/envs/examples/airl_envs/twod_maze.py
+++ /dev/null
@@ -1,62 +0,0 @@
-import numpy as np
-from gym import utils
-
-from imitation.envs.examples.airl_envs.common import get_asset_xml
-from imitation.envs.examples.airl_envs.twod_mjc_env import TwoDEnv
-
-INIT_POS = np.array([0.15, 0.15])
-TARGET = np.array([0.15, -0.15])
-DIST_THRESH = 0.12
-
-
-class TwoDMaze(TwoDEnv, utils.EzPickle):
-    def __init__(self, verbose=False):
-        self.verbose = verbose
-        self.max_episode_length = 200
-        self.episode_length = 0
-        utils.EzPickle.__init__(self)
-        TwoDEnv.__init__(
-            self,
-            get_asset_xml("twod_maze.xml"),
-            2,
-            xbounds=[-0.3, 0.3],
-            ybounds=[-0.3, 0.3],
-        )
-
-    def step(self, a):
-        self.do_simulation(a, self.frame_skip)
-        ob = self._get_obs()
-        pos = ob[0:2]
-        dist = np.sum(np.abs(pos - TARGET))  # np.linalg.norm(pos - TARGET)
-        reward = -(dist)
-
-        reward_ctrl = -np.square(a).sum()
-        reward += 1e-3 * reward_ctrl
-
-        if self.verbose:
-            print(pos, reward)
-        self.episode_length += 1
-        done = self.episode_length >= self.max_episode_length
-        return ob, reward, done, {"distance": dist}
-
-    def reset_model(self):
-        self.episode_length = 0
-        qpos = self.init_qpos + self.np_random.uniform(
-            size=self.model.nq,
-            low=-0.01,
-            high=0.01,
-        )
-        qvel = self.init_qvel + self.np_random.uniform(
-            size=self.model.nv,
-            low=-0.01,
-            high=0.01,
-        )
-        self.set_state(qpos, qvel)
-        return self._get_obs()
-
-    def _get_obs(self):
-        # return np.concatenate([self.sim.data.qpos, self.sim.data.qvel]).ravel()
-        return np.concatenate([self.sim.data.qpos]).ravel() - INIT_POS
-
-    def viewer_setup(self):
-        pass
diff --git a/src/imitation/envs/examples/airl_envs/twod_mjc_env.py b/src/imitation/envs/examples/airl_envs/twod_mjc_env.py
deleted file mode 100644
index d1f98b911..000000000
--- a/src/imitation/envs/examples/airl_envs/twod_mjc_env.py
+++ /dev/null
@@ -1,96 +0,0 @@
-from typing import NamedTuple, Tuple
-
-import numpy as np
-from gym.envs.mujoco import mujoco_env
-from gym.spaces import Box
-
-
-class MapConfig(NamedTuple):
-    xs: Tuple[float, float] = (-0.3, 0.3)
-    ys: Tuple[float, float] = (-0.3, 0.3)
-    xres: int = 50
-    yres: int = 50
-
-
-def make_heat_map(eval_func, map_config):
-    gps = get_dense_gridpoints(map_config)
-    vals = np.zeros(map_config.xres * map_config.yres)
-    for i, pnt in enumerate(gps):
-        vals[i] = eval_func(pnt)
-    return predictions_to_heatmap(vals, map_config)
-
-
-def get_dense_gridpoints(map_config):
-    xl = np.linspace(map_config.xs[0], map_config.xs[1], num=map_config.xres)
-    yl = np.linspace(map_config.ys[0], map_config.ys[1], num=map_config.yres)
-    gridpoints = np.zeros((map_config.xres * map_config.yres, 2))
-    for i in range(map_config.xres):
-        for j in range(map_config.yres):
-            gridpoints[i + map_config.xres * j] = np.array((xl[i], yl[j]))
-    return gridpoints
-
-
-def predictions_to_heatmap(predictions, map_config):
-    res = np.zeros((map_config.xres, map_config.yres))
-    for i in range(map_config.xres):
-        for j in range(map_config.yres):
-            res[i, j] = predictions[i + map_config.xres * j]
-    res = res / np.max(res)
-    return res.T
-
-
-def make_density_map(paths, map_config):
-    xs = np.linspace(map_config.xs[0], map_config.xs[1], num=map_config.xres + 1)
-    ys = np.linspace(map_config.ys[0], map_config.ys[1], num=map_config.yres + 1)
-    y = paths[:, 0]
-    x = paths[:, 1]
-    H, xedges, yedges = np.histogram2d(y, x, bins=(xs, ys))
-    H = H.astype(np.float)
-    H = H / np.max(H)
-    return H.T
-
-
-def plot_maps(combined_list=None, *heatmaps):
-    import matplotlib.pyplot as plt
-
-    combined = np.c_[heatmaps]
-    if combined_list is not None:
-        combined_list.append(combined)
-        combined = np.concatenate(combined_list)
-    else:
-        combined_list = []
-    plt.figure()
-    plt.imshow(combined, cmap="afmhot", interpolation="none")
-    plt.show()
-    return combined_list
-
-
-class TwoDEnv(mujoco_env.MujocoEnv):
-    def __init__(self, model_path, frame_skip, xbounds, ybounds):
-        super(TwoDEnv, self).__init__(model_path=model_path, frame_skip=frame_skip)
-        assert isinstance(self.observation_space, Box)
-        assert self.observation_space.shape == (2,)
-        self.__map_config = MapConfig(
-            xs=(xbounds[0], xbounds[1]),
-            ys=(ybounds[0], ybounds[1]),
-        )
-
-    @property
-    def map_config(self):
-        return self.__map_config
-
-    @property
-    def grid_flat_dim(self):
-        return self.map_config.xres * self.map_config.yres
-
-    def make_density_map(self, paths):
-        return make_density_map(paths, self.map_config)
-
-    def make_heatmap(self, eval_func):
-        return make_heat_map(eval_func, self.map_config)
-
-    def get_dense_gridpoints(self):
-        return get_dense_gridpoints(self.map_config)
-
-    def predictions_to_heatmap(self, predictions):
-        return predictions_to_heatmap(predictions, self.map_config)
diff --git a/src/imitation/envs/examples/airl_envs/utils.py b/src/imitation/envs/examples/airl_envs/utils.py
deleted file mode 100644
index b640d920c..000000000
--- a/src/imitation/envs/examples/airl_envs/utils.py
+++ /dev/null
@@ -1,37 +0,0 @@
-import numpy as np
-
-
-def flat_to_one_hot(val, ndim):
-    """
-
-    >>> flat_to_one_hot(2, ndim=4)
-    array([ 0.,  0.,  1.,  0.])
-    >>> flat_to_one_hot(4, ndim=5)
-    array([ 0.,  0.,  0.,  0.,  1.])
-    >>> flat_to_one_hot(np.array([2, 4, 3]), ndim=5)
-    array([[ 0.,  0.,  1.,  0.,  0.],
-           [ 0.,  0.,  0.,  0.,  1.],
-           [ 0.,  0.,  0.,  1.,  0.]])
-    """
-    shape = np.array(val).shape
-    v = np.zeros(shape + (ndim,))
-    if len(shape) == 1:
-        v[np.arange(shape[0]), val] = 1.0
-    else:
-        v[val] = 1.0
-    return v
-
-
-def one_hot_to_flat(val):
-    """
-    >>> one_hot_to_flat(np.array([0,0,0,0,1]))
-    4
-    >>> one_hot_to_flat(np.array([0,0,1,0]))
-    2
-    >>> one_hot_to_flat(np.array([[0,0,1,0], [1,0,0,0], [0,1,0,0]]))
-    array([2, 0, 1])
-    """
-    idxs = np.array(np.where(val == 1.0))[-1]
-    if len(val.shape) == 1:
-        return int(idxs)
-    return idxs
diff --git a/src/imitation/scripts/config/expert_demos.py b/src/imitation/scripts/config/expert_demos.py
index 03205432e..613e32a0a 100644
--- a/src/imitation/scripts/config/expert_demos.py
+++ b/src/imitation/scripts/config/expert_demos.py
@@ -132,6 +132,11 @@ def reacher():
     env_name = "Reacher-v2"
 
 
+@expert_demos_ex.named_config
+def seals_ant():
+    env_name = "seals/Ant-v0"
+
+
 @expert_demos_ex.named_config
 def seals_swimmer():
     env_name = "seals/Swimmer-v0"
@@ -142,26 +147,6 @@ def seals_walker():
     env_name = "seals/Walker2d-v0"
 
 
-# Custom env configs
-
-
-@expert_demos_ex.named_config
-def custom_ant():
-    env_name = "imitation/CustomAnt-v0"
-    locals().update(**ant_shared_locals)
-
-
-@expert_demos_ex.named_config
-def disabled_ant():
-    env_name = "imitation/DisabledAnt-v0"
-    locals().update(**ant_shared_locals)
-
-
-@expert_demos_ex.named_config
-def two_d_maze():
-    env_name = "imitation/TwoDMaze-v0"
-
-
 # Debug configs
 
 
diff --git a/src/imitation/scripts/config/train_adversarial.py b/src/imitation/scripts/config/train_adversarial.py
index 2dca3065d..a0eeda07b 100644
--- a/src/imitation/scripts/config/train_adversarial.py
+++ b/src/imitation/scripts/config/train_adversarial.py
@@ -256,35 +256,6 @@ def seals_walker():
     rollout_hint = "walker"
 
 
-# Custom Gym environment named configs
-
-
-@train_adversarial_ex.named_config
-def two_d_maze():
-    locals().update(**MUJOCO_SHARED_LOCALS)
-    env_name = "imitation/TwoDMaze-v0"
-    rollout_hint = "two_d_maze"
-
-
-@train_adversarial_ex.named_config
-def custom_ant():
-    locals().update(**MUJOCO_SHARED_LOCALS)
-    # Watch out -- ANT_SHARED_LOCALS could erroneously erase nested dict keys from
-    # MUJOCO_SHARED_LOCALS because `locals().update()` doesn't merge dicts
-    # "Sacred-style".
-    locals().update(**ANT_SHARED_LOCALS)
-    env_name = "imitation/CustomAnt-v0"
-    rollout_hint = "custom_ant"
-
-
-@train_adversarial_ex.named_config
-def disabled_ant():
-    locals().update(**MUJOCO_SHARED_LOCALS)
-    locals().update(**ANT_SHARED_LOCALS)
-    env_name = "imitation/DisabledAnt-v0"
-    rollout_hint = "disabled_ant"
-
-
 # Debug configs
 
 


Test patch:
diff --git a/tests/test_scripts.py b/tests/test_scripts.py
index 621f7cb5c..8eb659f26 100644
--- a/tests/test_scripts.py
+++ b/tests/test_scripts.py
@@ -470,10 +470,10 @@ def test_parallel_train_adversarial_custom_env(tmpdir):
     import gym
 
     try:
-        gym.make("Ant-v3")
+        gym.make("seals/Ant-v0")
     except gym.error.DependencyNotInstalled:  # pragma: no cover
         pytest.skip("mujoco_py not available")
-    env_named_config = "custom_ant"
+    env_named_config = "seals_ant"
     rollout_path = _generate_test_rollouts(tmpdir, env_named_config)
 
     config_updates = dict(