sim.py

import os
import queue
import threading
from dataclasses import dataclass, field

import mujoco
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
FPS = 30
PHYSICS_DT = 1.0 / 600.0
SUBSTEPS = int(round(1.0 / FPS / PHYSICS_DT))
JOINTS = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
GRIPPER_CLOSED = -0.17453
GRIPPER_OPEN = 1.74533
REST_DEG = np.array([0.0, -92.0, 88.0, 72.0, 0.0, 2.0])

COLORS = {
    "red": "0.80 0.12 0.10 1",
    "blue": "0.10 0.25 0.75 1",
    "green": "0.12 0.55 0.20 1",
    "yellow": "0.95 0.78 0.10 1",
    "white": "0.92 0.92 0.90 1",
    "black": "0.08 0.08 0.08 1",
    "orange": "0.95 0.45 0.08 1",
}

SCENES = {
    "Cubes and tray": {
        "objects": [
            ("red", "cube", (0.20, 0.0), 0.0),
            ("blue", "cube", (0.25, 0.09), 0.4),
            ("yellow", "cube", (0.16, 0.11), 0.9),
        ],
        "tray": (0.24, -0.16),
    },
    "Single cube": {
        "objects": [("red", "cube", (0.20, -0.01), 0.3)],
        "tray": (0.24, -0.16),
    },
    "Blocks without tray": {
        "objects": [
            ("green", "cube", (0.22, 0.08), 0.2),
            ("orange", "cube", (0.24, -0.06), 0.7),
            ("white", "box", (0.17, 0.0), 0.0),
        ],
        "tray": None,
    },
}


def _object_xml(i, color, kind, xy, yaw):
    rgba = COLORS[color]
    half = "0.015 0.015 0.015" if kind == "cube" else "0.02 0.03 0.015"
    z = 0.0155
    q = f"{np.cos(yaw / 2):.5f} 0 0 {np.sin(yaw / 2):.5f}"
    return (
        f'<body name="obj{i}" pos="{xy[0]} {xy[1]} {z}" quat="{q}">'
        f'<freejoint name="obj{i}_free"/>'
        f'<geom type="box" size="{half}" rgba="{rgba}" mass="0.02" friction="1.5 0.02 0.001" '
        f'condim="4" solref="0.004 1" priority="1"/>'
        f"</body>"
    )


def _tray_xml(xy):
    x, y = xy
    hx, hy, hz, t = 0.075, 0.06, 0.022, 0.004
    rgba = "0.55 0.56 0.58 1"
    parts = [
        f'<geom type="box" pos="0 0 {t / 2}" size="{hx} {hy} {t / 2}" rgba="{rgba}"/>',
        f'<geom type="box" pos="{hx - t / 2} 0 {hz}" size="{t / 2} {hy} {hz}" rgba="{rgba}"/>',
        f'<geom type="box" pos="{-hx + t / 2} 0 {hz}" size="{t / 2} {hy} {hz}" rgba="{rgba}"/>',
        f'<geom type="box" pos="0 {hy - t / 2} {hz}" size="{hx} {t / 2} {hz}" rgba="{rgba}"/>',
        f'<geom type="box" pos="0 {-hy + t / 2} {hz}" size="{hx} {t / 2} {hz}" rgba="{rgba}"/>',
    ]
    return f'<body name="tray" pos="{x} {y} 0">' + "".join(parts) + "</body>"


def scene_xml(scene):
    spec = SCENES[scene]
    objects = "".join(_object_xml(i, *o) for i, o in enumerate(spec["objects"]))
    tray = _tray_xml(spec["tray"]) if spec["tray"] else ""
    return f"""
<mujoco model="so101_tabletop">
  <include file="so101.xml"/>
  <option timestep="{PHYSICS_DT}" cone="elliptic" impratio="10"/>
  <visual>
    <headlight diffuse="0.4 0.4 0.4" ambient="0.5 0.5 0.5" specular="0 0 0"/>
    <quality shadowsize="1024"/>
    <global offwidth="640" offheight="480"/>
    <map znear="0.005"/>
  </visual>
  <asset>
    <texture type="skybox" builtin="gradient" rgb1="0.85 0.86 0.88" rgb2="0.55 0.56 0.58" width="256" height="1536"/>
    <texture name="cork" type="2d" builtin="flat" rgb1="0.80 0.66 0.48" mark="random" markrgb="0.70 0.55 0.38" random="0.12" width="512" height="512"/>
    <material name="cork" texture="cork" texrepeat="4 4" reflectance="0.0"/>
    <material name="desk" rgba="0.80 0.81 0.82 1"/>
    <material name="wall" rgba="0.90 0.90 0.88 1"/>
    <material name="white_part" rgba="0.93 0.93 0.91 1"/>
  </asset>
  <worldbody>
    <light pos="0.3 -0.3 1.2" dir="-0.2 0.2 -1" directional="false" diffuse="0.35 0.35 0.35" castshadow="true"/>
    <light pos="-0.4 0.4 1.0" dir="0.3 -0.3 -1" diffuse="0.3 0.3 0.3" castshadow="false"/>
    <geom name="desk" type="box" pos="0.25 0 -0.012" size="0.7 0.7 0.004" material="desk"/>
    <geom name="board" type="box" pos="0.22 0 -0.004" size="0.26 0.32 0.004" material="cork" friction="1 0.01 0.001"/>
    <geom name="wall" type="box" pos="-0.35 0 0.4" size="0.01 1.2 0.5" material="wall" contype="0" conaffinity="0"/>
    <geom name="wall2" type="box" pos="0.3 0.75 0.4" size="1.0 0.01 0.5" material="wall" contype="0" conaffinity="0"/>
    <camera name="scene" pos="0.30 0.37 0.35" mode="targetbody" target="scene_target" fovy="58"/>
    <camera name="overview" pos="0.40 0.52 0.46" mode="targetbody" target="scene_target" fovy="50"/>
    <body name="scene_target" pos="0.17 -0.04 0.02"/>
    {objects}
    {tray}
  </worldbody>
</mujoco>
"""


def build_model(scene):
    xml = scene_xml(scene)
    cwd = os.getcwd()
    os.chdir(HERE)
    try:
        model = mujoco.MjModel.from_xml_string(xml)
    finally:
        os.chdir(cwd)
    return model


def deg_to_q(deg):
    q = np.deg2rad(np.asarray(deg[:5], dtype=np.float64))
    g = GRIPPER_CLOSED + np.clip(deg[5], 0.0, 100.0) / 100.0 * (GRIPPER_OPEN - GRIPPER_CLOSED)
    return np.concatenate([q, [g]])


def q_to_deg(q):
    deg = np.rad2deg(np.asarray(q[:5], dtype=np.float64))
    g = (q[5] - GRIPPER_CLOSED) / (GRIPPER_OPEN - GRIPPER_CLOSED) * 100.0
    return np.concatenate([deg, [g]])


@dataclass
class Sim:
    scene: str = "Cubes and tray"
    width: int = 256
    height: int = 256
    seed: int = 0
    offscreen: bool = True
    model: mujoco.MjModel = field(init=False)
    data: mujoco.MjData = field(init=False)

    def __post_init__(self):
        self.model = build_model(self.scene)
        self.data = mujoco.MjData(self.model)
        self.qadr = np.array([self.model.joint(j).qposadr[0] for j in JOINTS])
        self.renderer = mujoco.Renderer(self.model, self.height, self.width) if self.offscreen else None
        self.reset()

    def reset(self):
        mujoco.mj_resetData(self.model, self.data)
        rng = np.random.default_rng(self.seed)
        for i in range(self.model.nbody):
            name = self.model.body(i).name
            if name.startswith("obj"):
                j = self.model.joint(f"{name}_free")
                a = j.qposadr[0]
                self.data.qpos[a : a + 2] += rng.uniform(-0.015, 0.015, 2)
        q = deg_to_q(REST_DEG)
        self.data.qpos[self.qadr] = q
        self.data.ctrl[:] = q
        mujoco.mj_forward(self.model, self.data)
        for _ in range(120):
            mujoco.mj_step(self.model, self.data)

    def state_deg(self):
        return q_to_deg(self.data.qpos[self.qadr])

    def step(self, command_deg):
        q = deg_to_q(command_deg)
        lo = self.model.actuator_ctrlrange[:, 0]
        hi = self.model.actuator_ctrlrange[:, 1]
        self.data.ctrl[:] = np.clip(q, lo, hi)
        for _ in range(SUBSTEPS):
            mujoco.mj_step(self.model, self.data)

    def render(self, camera):
        self.renderer.update_scene(self.data, camera=camera)
        return self.renderer.render().copy()

    def close(self):
        if self.renderer is not None:
            self.renderer.close()

    def object_positions(self):
        out = {}
        for i in range(self.model.nbody):
            name = self.model.body(i).name
            if name.startswith("obj"):
                out[name] = self.data.body(i).xpos.copy()
        return out


class Recorder(threading.Thread):
    def __init__(self, model, path, fps, width=480, height=360, camera="overview"):
        super().__init__(daemon=True)
        self.model = model
        self.path = path
        self.fps = fps
        self.size = (width, height)
        self.camera = camera
        self.items = queue.Queue()
        self.busy = 0.0
        self.error = None

    def push(self, qpos, caption):
        self.items.put((qpos.copy(), caption))

    def finish(self):
        self.items.put(None)
        self.join()
        return self.path if self.error is None else None

    def run(self):
        import time

        import imageio.v2 as imageio
        from PIL import Image, ImageDraw

        data = mujoco.MjData(self.model)
        renderer = mujoco.Renderer(self.model, self.size[1], self.size[0])
        writer = imageio.get_writer(self.path, fps=self.fps, codec="libx264", quality=7, macro_block_size=8)
        try:
            while True:
                item = self.items.get()
                if item is None:
                    break
                t0 = time.perf_counter()
                data.qpos[:] = item[0]
                mujoco.mj_forward(self.model, data)
                renderer.update_scene(data, camera=self.camera)
                im = Image.fromarray(renderer.render())
                d = ImageDraw.Draw(im)
                d.rectangle([0, self.size[1] - 22, self.size[0], self.size[1]], fill=(0, 0, 0))
                d.text((6, self.size[1] - 17), item[1], fill=(255, 255, 255))
                writer.append_data(np.asarray(im))
                self.busy += time.perf_counter() - t0
        except Exception as e:
            self.error = e
        finally:
            writer.close()
            renderer.close()