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''
f''
f''
f""
)
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'',
f'',
f'',
f'',
f'',
]
return f'' + "".join(parts) + ""
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"""
{objects}
{tray}
"""
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()