Buckets:

twanghcmut
/

ego-real2sim

Files

xet

twanghcmut/ego-real2sim / basic_pick_place_clean /code /src /real2sim /sim /physics.py

twanghcmut

20 days ago

download

raw

27.8 kB

import json
import shutil
import tempfile
import xml.etree.ElementTree as ET
from pathlib import Path

import mujoco
import numpy as np
from PIL import Image
from scipy.spatial import ConvexHull
from scipy.spatial.transform import Rotation

from ..geometry import MUJOCO_CAMERA_FROM_OPENCV, quat_wxyz


class PhysicsScene:
    def __init__(self, folder: Path, asset: Path, reference: Path | None = None):
        self.folder = Path(folder).resolve()
        self.asset = Path(asset).resolve()
        self._temporary = tempfile.TemporaryDirectory(prefix="real2sim-physics-")
        self._renderers = {}
        with np.load(self.folder / "sim_trajectory.npz") as trajectory:
            self.reference_qpos = trajectory["qpos"].copy()
            self.camera = trajectory["camera"].copy()
            self.timestamps = trajectory["timestamps"].copy()
            self.frame_indices = trajectory["frame_indices"].copy()
        if reference is not None:
            with np.load(reference) as fitted:
                if fitted["qpos"].shape != self.reference_qpos.shape:
                    raise ValueError("Fitted hand trajectory must match the reference")
                if "timestamps" in fitted and not np.allclose(
                    fitted["timestamps"], self.timestamps
                ):
                    raise ValueError("Fitted hand timestamps must match the reference")
                self.reference_qpos = fitted["qpos"].copy()
        with np.load(self.folder / "real_masks.npz") as target:
            names = target["names"].tolist()
            self.real_masks = target["masks"][
                [names.index(n) for n in ("left", "right", "object")]
            ].copy()
        with np.load(self.asset) as geometry:
            vertices = np.asarray(geometry["world_vertices"], dtype=float)
            self.object_faces = np.asarray(geometry["faces"], dtype=np.int32)
            self.object_colors = np.asarray(geometry["colors"], dtype=np.uint8)
            self.source_frame_index = int(geometry["source_frame_index"])
        if (
            vertices.ndim != 2
            or vertices.shape[1] != 3
            or not np.isfinite(vertices).all()
        ):
            raise ValueError("Asset must contain finite world vertices")
        if self.object_faces.ndim != 2 or self.object_faces.shape[1] != 3:
            raise ValueError("Asset must contain triangular faces")
        if self.object_colors.shape != vertices.shape:
            raise ValueError("Asset colors must match vertices")
        if len(self.timestamps) < 2 or np.any(np.diff(self.timestamps) <= 0):
            raise ValueError("Trajectory timestamps must increase")
        self.object_center = (vertices.min(axis=0) + vertices.max(axis=0)) / 2
        self.local_vertices = vertices - self.object_center
        self._prepare_table()
        self.model = self._compile()
        self.data = mujoco.MjData(self.model)
        self.object_joint = self.model.joint("object_free").id
        self.object_qpos_address = int(self.model.jnt_qposadr[self.object_joint])
        self.object_dof_address = int(self.model.jnt_dofadr[self.object_joint])
        self.object_geom = self.model.geom("object_collision").id
        self.object_visual_geom = self.model.geom("object_visual").id
        self.camera_mocap = int(
            self.model.body_mocapid[self.model.body("ego_mount").id]
        )
        self.hand_geoms = []
        self.hand_joints = []
        self.hand_visible = []
        self.hand_mocap = []
        self.hand_equalities = []
        self.wrist_addresses = []
        self.ball_joints = []
        for side in ("left", "right"):
            joints = np.array(
                [
                    j
                    for j in range(self.model.njnt)
                    if self.model.joint(j).name.startswith(side + "_")
                ]
            )
            geoms = np.array(
                [
                    g
                    for g in range(self.model.ngeom)
                    if self.model.body(int(self.model.geom_bodyid[g])).name.startswith(
                        side + "_"
                    )
                ]
            )
            root = next(
                j
                for j in joints
                if self.model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE
            )
            address = int(self.model.jnt_qposadr[root])
            self.hand_joints.append(joints)
            self.hand_geoms.append(geoms)
            self.wrist_addresses.append(address)
            self.hand_visible.append(self.reference_qpos[:, address + 2] > -9)
            self.hand_mocap.append(
                int(self.model.body_mocapid[self.model.body(side + "_target").id])
            )
            self.hand_equalities.append(self.model.equality(side + "_weld").id)
            self.ball_joints.extend(
                j
                for j in joints
                if self.model.jnt_type[j] == mujoco.mjtJoint.mjJNT_BALL
            )
        self.ball_joints = np.asarray(self.ball_joints, dtype=int)
        self.actor_names = np.array(["left", "right", "object"])
        self.geom_actor_ids = np.full(self.model.ngeom, -1, dtype=np.int8)
        for actor, geoms in enumerate(self.hand_geoms):
            self.geom_actor_ids[geoms] = actor
        self.geom_actor_ids[[self.object_geom, self.object_visual_geom]] = 2
        self.ball_addresses = self.model.jnt_qposadr[self.ball_joints]
        self._finger_targets = np.stack(
            [
                Rotation.from_quat(
                    self.reference_qpos[:, a : a + 4][:, [1, 2, 3, 0]]
                ).as_rotvec()
                for a in self.ball_addresses
            ],
            axis=1,
        ).reshape(len(self.timestamps), -1)
        self.parameter_mean = np.r_[np.zeros(24), 0.8]
        self.parameter_std = np.r_[
            np.full(3, 0.015), np.full(3, 0.12), np.full(18, 0.015), 0.2
        ]
        self.parameter_bounds = np.column_stack(
            (
                np.r_[np.full(3, -0.06), np.full(3, -0.5), np.full(18, -0.08), 0.15],
                np.r_[np.full(3, 0.06), np.full(3, 0.5), np.full(18, 0.08), 1.8],
            )
        )
        self.bounds = self.parameter_bounds
        self.parameter_names = [
            "object_" + n for n in ("x", "y", "z", "rx", "ry", "rz")
        ]
        self.parameter_names += [
            f"{side}_{key}_{axis}"
            for side in ("left", "right")
            for key in range(3)
            for axis in ("x", "y", "z")
        ]
        self.parameter_names += ["friction"]
        self.reset(self.parameter_mean)

    def _prepare_table(self):
        scene = json.loads((self.folder / "scene.json").read_text())
        if scene.get("geometry_source") != "flat_table":
            raise ValueError("Scene geometry_source must be flat_table")
        height = scene.get("table_height_m")
        if (
            isinstance(height, bool)
            or not isinstance(height, (int, float))
            or not np.isfinite(height)
        ):
            raise ValueError("Scene table_height_m must be a finite number")
        source = scene.get("table_source")
        if not isinstance(source, str) or not source.strip():
            raise ValueError("Scene table_source must identify the height estimate")
        self.table_height = float(height)
        self.table_diagnostics = {
            "height_m": self.table_height,
            "source": source,
            "geometry_source": "flat_table",
        }

    def _compile(self):
        spec = mujoco.MjSpec.from_file(str(self.folder / "scene.xml"))
        spec.option.timestep = 1 / 600
        spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
        spec.option.iterations = 80
        spec.option.gravity = [0, 0, -9.81]
        spec.visual.global_.offwidth = 1920
        spec.visual.global_.offheight = 1080
        spec.delete(spec.geom("floor"))
        for material in list(spec.materials):
            if material.name == "grid" and not any(
                geom.material == material.name for geom in spec.geoms
            ):
                spec.delete(material)
        for texture in list(spec.textures):
            if texture.name == "grid" and not any(
                texture.name in list(material.textures) for material in spec.materials
            ):
                spec.delete(texture)
        spec.visual.headlight.ambient = [0.35, 0.35, 0.35]
        spec.visual.headlight.diffuse = [0.55, 0.55, 0.55]
        spec.visual.headlight.specular = [0.05, 0.05, 0.05]
        spec.visual.quality.shadowsize = 2048
        for light in list(spec.lights):
            spec.delete(light)
        spec.worldbody.add_light(
            name="key_light",
            type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL,
            pos=self.object_center + np.array([0.5, -0.5, 2.0]),
            dir=[-0.25, 0.25, -1],
            ambient=[0.1, 0.1, 0.1],
            diffuse=[0.6, 0.6, 0.6],
            specular=[0.05, 0.05, 0.05],
            castshadow=True,
        )
        spec.add_texture(
            name="background",
            type=mujoco.mjtTexture.mjTEXTURE_SKYBOX,
            builtin=mujoco.mjtBuiltin.mjBUILTIN_FLAT,
            rgb1=[0.88, 0.89, 0.91],
            rgb2=[0.88, 0.89, 0.91],
            width=128,
            height=768,
        )
        spec.add_material(
            name="table_material",
            rgba=[0.78, 0.74, 0.65, 1],
            specular=0,
            shininess=0,
            reflectance=0,
        )
        spec.worldbody.add_geom(
            name="table",
            type=mujoco.mjtGeom.mjGEOM_PLANE,
            pos=[0, 0, self.table_height],
            size=[2, 2, 0.1],
            contype=1,
            conaffinity=6,
            material="table_material",
            friction=[0.8, 0.01, 0.001],
        )
        width = int(np.ceil(np.sqrt(len(self.object_faces))))
        height = int(np.ceil(len(self.object_faces) / width))
        pixels = np.zeros((height * width, 3), dtype=np.uint8)
        pixels[: len(self.object_faces)] = (
            self.object_colors[self.object_faces].mean(axis=1).astype(np.uint8)
        )
        self.object_texture = Path(self._temporary.name) / "object_colors.png"
        Image.fromarray(pixels.reshape(height, width, 3)).save(self.object_texture)
        spec.add_texture(
            name="object_colors",
            type=mujoco.mjtTexture.mjTEXTURE_2D,
            file=str(self.object_texture),
        )
        material = spec.add_material(
            name="object_material", texuniform=False, specular=0
        )
        material.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "object_colors"
        atlas = np.arange(len(self.object_faces))
        self.object_uv = np.column_stack(
            ((atlas % width + 0.5) / width, (atlas // width + 0.5) / height)
        )
        spec.add_mesh(
            name="object_mesh",
            uservert=self.local_vertices.ravel(),
            userface=self.object_faces.ravel(),
            usertexcoord=self.object_uv.ravel(),
            userfacetexcoord=np.repeat(atlas, 3),
        )
        hull = ConvexHull(self.local_vertices)
        used, faces = np.unique(hull.simplices, return_inverse=True)
        spec.add_mesh(
            name="object_hull",
            uservert=self.local_vertices[used].ravel(),
            userface=faces.ravel(),
        )
        body = spec.worldbody.add_body(name="object", pos=self.object_center)
        body.add_joint(name="object_free", type=mujoco.mjtJoint.mjJNT_FREE)
        body.add_geom(
            name="object_visual",
            type=mujoco.mjtGeom.mjGEOM_MESH,
            meshname="object_mesh",
            material="object_material",
            contype=0,
            conaffinity=0,
            mass=0,
        )
        body.add_geom(
            name="object_collision",
            type=mujoco.mjtGeom.mjGEOM_MESH,
            meshname="object_hull",
            group=3,
            rgba=[0, 0, 0, 0],
            mass=0.1,
            contype=4,
            conaffinity=3,
            condim=4,
            priority=1,
            friction=[0.8, 0.01, 0.001],
            solref=[0.01, 1],
        )
        for geom in spec.geoms:
            if geom.parent.name.startswith(("left_", "right_")):
                geom.contype = 2
                geom.conaffinity = 5
                geom.friction = [0.8, 0.01, 0.001]
        for side in ("left", "right"):
            spec.worldbody.add_body(name=side + "_target", mocap=True)
            spec.add_equality(
                name=side + "_weld",
                type=mujoco.mjtEq.mjEQ_WELD,
                name1=side + "_target",
                name2=side + "_Wrist",
                objtype=mujoco.mjtObj.mjOBJ_BODY,
                data=[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1],
                solref=[0.015, 1],
            )
        for joint in spec.joints:
            if joint.type != mujoco.mjtJoint.mjJNT_BALL:
                continue
            joint.damping = [0.015, 0, 0]
            joint.armature = 0.0001
            for axis in range(3):
                gear = np.zeros(6)
                gear[axis] = 1
                spec.add_actuator(
                    name=f"{joint.name}_{axis}",
                    trntype=mujoco.mjtTrn.mjTRN_JOINT,
                    target=joint.name,
                    gear=gear,
                    gaintype=mujoco.mjtGain.mjGAIN_FIXED,
                    gainprm=[0.8] + [0] * 9,
                    biastype=mujoco.mjtBias.mjBIAS_AFFINE,
                    biasprm=[0, -0.8, -0.025] + [0] * 7,
                    ctrllimited=True,
                    ctrlrange=[-np.pi, np.pi],
                    forcelimited=True,
                    forcerange=[-2, 2],
                )
        self.spec = spec
        return spec.compile()

    def _parameters(self, parameters):
        values = np.asarray(parameters, dtype=float)
        if values.shape != (25,) or not np.isfinite(values).all():
            raise ValueError("Expected 25 finite parameters")
        if np.any(values < self.parameter_bounds[:, 0]) or np.any(
            values > self.parameter_bounds[:, 1]
        ):
            raise ValueError("Parameters exceed bounds")
        return values.copy()

    def reset(self, parameters):
        self.parameters = self._parameters(parameters)
        mujoco.mj_resetData(self.model, self.data)
        self.data.qpos[: self.reference_qpos.shape[1]] = self.reference_qpos[0]
        rotation = Rotation.from_rotvec(self.parameters[3:6])
        quaternion = rotation.as_quat()[[3, 0, 1, 2]]
        position = self.object_center + self.parameters[:3]
        minimum_z = float((self.local_vertices @ rotation.as_matrix().T)[:, 2].min())
        self.initial_table_adjustment = max(
            0.0, self.table_height + 0.001 - (position[2] + minimum_z)
        )
        position[2] += self.initial_table_adjustment
        address = self.object_qpos_address
        self.data.qpos[address : address + 7] = np.r_[position, quaternion]
        self.model.geom_friction[self.object_geom, 0] = self.parameters[-1]
        self._active = np.zeros(2, dtype=bool)
        self._contact_start = np.full(2, np.inf)
        self._activation_frames = [[], []]
        self._collision_enabled = np.zeros(2, dtype=bool)
        self._contact_steps = 0
        self._contact_pairs = 0
        self._contact_peak_force = 0.0
        self._contact_impulse = 0.0
        self._contact_by_hand = np.zeros(2, dtype=int)
        self._initial_object_position = position.copy()
        self._maximum_object_z = position[2]
        self._unstable = False
        for geoms in self.hand_geoms:
            self.model.geom_contype[geoms] = 0
            self.model.geom_conaffinity[geoms] = 0
        self._drive(0.0)
        mujoco.mj_forward(self.model, self.data)

    def _drive(self, time):
        high = min(
            int(np.searchsorted(self.timestamps, time, side="right")),
            len(self.timestamps) - 1,
        )
        low = max(0, high - 1)
        fraction = np.clip(
            (time - self.timestamps[low])
            / (self.timestamps[high] - self.timestamps[low]),
            0,
            1,
        )
        nearest = min(
            int(np.searchsorted(self.timestamps, time + 1e-9, side="right")) - 1,
            len(self.timestamps) - 1,
        )
        nearest = max(0, nearest)
        control = (1 - fraction) * self._finger_targets[
            low
        ] + fraction * self._finger_targets[high]
        self.data.ctrl[:] = control
        corrections = self.parameters[6:24].reshape(2, 3, 3)
        key_times = np.linspace(self.timestamps[0], self.timestamps[-1], 3)
        for side, address in enumerate(self.wrist_addresses):
            visible = bool(self.hand_visible[side][nearest])
            mocap = self.hand_mocap[side]
            if visible:
                correction = np.array(
                    [
                        np.interp(time, key_times, corrections[side, :, axis])
                        for axis in range(3)
                    ]
                )
                a, b = self.reference_qpos[[low, high], address : address + 7].copy()
                if not self.hand_visible[side][low]:
                    a = b.copy()
                if not self.hand_visible[side][high]:
                    b = a.copy()
                position = (1 - fraction) * a[:3] + fraction * b[:3] + correction
                qb = b[3:] * (1 if np.dot(a[3:], b[3:]) >= 0 else -1)
                quaternion = (1 - fraction) * a[3:] + fraction * qb
                quaternion /= np.linalg.norm(quaternion)
                if not self._active[side]:
                    for joint in self.hand_joints[side]:
                        start = int(self.model.jnt_qposadr[joint])
                        size = (
                            7
                            if self.model.jnt_type[joint] == mujoco.mjtJoint.mjJNT_FREE
                            else 4
                        )
                        self.data.qpos[start : start + size] = self.reference_qpos[
                            nearest, start : start + size
                        ]
                        velocity = int(self.model.jnt_dofadr[joint])
                        count = 6 if size == 7 else 3
                        self.data.qvel[velocity : velocity + count] = 0
                    self.data.qpos[address : address + 7] = np.r_[position, quaternion]
                    self._contact_start[side] = time + 0.1
                    self._activation_frames[side].append(nearest)
                enabled = time >= self._contact_start[side]
            else:
                position = np.array([5.0 + side, 5.0, self.table_height + 1])
                quaternion = np.array([1.0, 0, 0, 0])
                enabled = False
                if self._active[side] or time == 0:
                    self.data.qpos[address : address + 7] = np.r_[position, quaternion]
                    root = int(self.model.jnt_dofadr[self.hand_joints[side][0]])
                    self.data.qvel[root : root + 6] = 0
            self.data.mocap_pos[mocap] = position
            self.data.mocap_quat[mocap] = quaternion
            if enabled != self._collision_enabled[side]:
                self.model.geom_contype[self.hand_geoms[side]] = 2 if enabled else 0
                self.model.geom_conaffinity[self.hand_geoms[side]] = 5 if enabled else 0
                self._collision_enabled[side] = enabled
            self._active[side] = visible

    def _contacts(self):
        force = np.zeros(6)
        touched = False
        for index in range(self.data.ncon):
            contact = self.data.contact[index]
            if contact.geom1 == self.object_geom:
                other = contact.geom2
            elif contact.geom2 == self.object_geom:
                other = contact.geom1
            else:
                continue
            for side, geoms in enumerate(self.hand_geoms):
                if other not in geoms:
                    continue
                mujoco.mj_contactForce(self.model, self.data, index, force)
                normal = max(0.0, float(force[0]))
                if normal <= 1e-8:
                    continue
                touched = True
                self._contact_pairs += 1
                self._contact_by_hand[side] += 1
                self._contact_peak_force = max(self._contact_peak_force, normal)
                self._contact_impulse += normal * self.model.opt.timestep
        self._contact_steps += int(touched)
        self._maximum_object_z = max(
            self._maximum_object_z, float(self.data.qpos[self.object_qpos_address + 2])
        )

    def _renderer(self, size):
        size = tuple(size)
        if size not in self._renderers:
            self._renderers[size] = mujoco.Renderer(
                self.model, width=size[0], height=size[1]
            )
        return self._renderers[size]

    def rollout(self, parameters, frame_indices, render_size=(256, 144), rgb=False):
        indices = np.asarray(frame_indices, dtype=int)
        if (
            indices.ndim != 1
            or not len(indices)
            or np.any(np.diff(indices) <= 0)
            or indices[0] < 0
            or indices[-1] >= len(self.timestamps)
        ):
            raise ValueError("Frame indices must increase within the saved trajectory")
        self.reset(parameters)
        renderer = self._renderer(render_size)
        masks, images, qpos, positions, full_qpos, counts, forces, geom_ids = (
            [],
            [],
            [],
            [],
            [],
            [],
            [],
            [],
        )
        for frame in indices:
            target_time = self.timestamps[frame] - self.timestamps[0]
            while self.data.time + self.model.opt.timestep / 2 < target_time:
                self._drive(self.data.time + self.timestamps[0])
                mujoco.mj_step(self.model, self.data)
                self._contacts()
                if (
                    not np.isfinite(self.data.qpos).all()
                    or np.max(np.abs(self.data.qvel)) > 1e4
                ):
                    self._unstable = True
                    raise FloatingPointError("Physics rollout became unstable")
            self._drive(self.data.time + self.timestamps[0])
            self.data.mocap_pos[self.camera_mocap] = self.camera[frame, :3, 3]
            self.data.mocap_quat[self.camera_mocap] = quat_wxyz(
                self.camera[frame, :3, :3] @ MUJOCO_CAMERA_FROM_OPENCV
            )
            mujoco.mj_forward(self.model, self.data)
            renderer.enable_segmentation_rendering()
            renderer.update_scene(self.data, camera="ego")
            segmentation = renderer.render()
            is_geom = segmentation[..., 1] == mujoco.mjtObj.mjOBJ_GEOM
            ids = segmentation[..., 0]
            geom_ids.append(np.where(is_geom, ids, -1).astype(np.int32))
            masks.append(
                np.stack(
                    [
                        is_geom & np.isin(ids, geoms)
                        for geoms in (*self.hand_geoms, [self.object_visual_geom])
                    ]
                )
            )
            if rgb:
                renderer.disable_segmentation_rendering()
                renderer.update_scene(self.data, camera="ego")
                images.append(renderer.render().copy())
            address = self.object_qpos_address
            qpos.append(self.data.qpos[address : address + 7].copy())
            positions.append(self.data.qpos[address : address + 3].copy())
            full_qpos.append(self.data.qpos.copy())
            counts.append(self._contact_pairs)
            forces.append(self._contact_peak_force)
        diagnostics = {
            "simulation_mode": "contact_physics",
            "hand_object_contact_steps": self._contact_steps,
            "hand_object_contact_pairs": self._contact_pairs,
            "hand_object_contacts_by_hand": self._contact_by_hand.tolist(),
            "peak_hand_object_normal_force_n": self._contact_peak_force,
            "hand_object_normal_impulse_ns": self._contact_impulse,
            "maximum_lift_m": float(
                self._maximum_object_z - self._initial_object_position[2]
            ),
            "final_displacement_m": float(
                np.linalg.norm(positions[-1] - self._initial_object_position)
            ),
            "initial_table_adjustment_m": float(self.initial_table_adjustment),
            "hand_activation_frames": self._activation_frames,
            "activation_collision_delay_s": 0.1,
            "table": self.table_diagnostics,
            "unstable": self._unstable,
        }
        result = {
            "masks": np.stack(masks, axis=1),
            "geom_ids": np.stack(geom_ids),
            "geom_actor_ids": self.geom_actor_ids.copy(),
            "actor_names": self.actor_names.copy(),
            "object_qpos": np.stack(qpos),
            "object_positions": np.stack(positions),
            "qpos": np.stack(full_qpos),
            "frame_indices": indices.copy(),
            "timestamps": self.timestamps[indices].copy(),
            "contact_counts": np.asarray(counts),
            "contact_forces": np.asarray(forces),
            "diagnostics": diagnostics,
        }
        if rgb:
            result["rgb"] = np.stack(images)
        return result

    def save(self, folder):
        folder = Path(folder)
        folder.mkdir(parents=True, exist_ok=True)
        mujoco.mj_saveModel(self.model, str(folder / "scene.mjb"), None)
        shutil.copyfile(self.object_texture, folder / "object_colors.png")
        with (folder / "object_visual.obj").open("w") as mesh:
            np.savetxt(mesh, self.local_vertices, fmt="v %.9g %.9g %.9g")
            np.savetxt(mesh, self.object_uv, fmt="vt %.9g %.9g")
            for index, face in enumerate(self.object_faces, start=1):
                mesh.write(
                    "f " + " ".join(f"{vertex + 1}/{index}" for vertex in face) + "\n"
                )
        xml = self.spec.to_xml().replace(str(self.object_texture), "object_colors.png")
        root = ET.fromstring(xml)
        for mesh in root.findall("asset/mesh"):
            if mesh.get("name") == "object_mesh":
                mesh.attrib.clear()
                mesh.set("name", "object_mesh")
                mesh.set("file", "object_visual.obj")
        for geom in root.findall(".//geom"):
            if geom.get("name") == "object_collision":
                geom.set("friction", f"{self.parameters[-1]:.9g} 0.01 0.001")
        (folder / "scene.xml").write_text(ET.tostring(root, encoding="unicode"))
        np.savez_compressed(
            folder / "controller.npz",
            reference_qpos=self.reference_qpos,
            camera=self.camera,
            timestamps=self.timestamps,
            parameters=self.parameters,
            parameter_mean=self.parameter_mean,
            parameter_std=self.parameter_std,
            parameter_bounds=self.parameter_bounds,
            initial_object_position=self._initial_object_position,
        )
        (folder / "scene_metrics.json").write_text(
            json.dumps(
                {
                    "table": self.table_diagnostics,
                    "mass_kg": 0.1,
                    "source_frame_index": self.source_frame_index,
                    "simulation_mode": "contact_physics",
                },
                indent=2,
            )
        )

    def close(self):
        for renderer in self._renderers.values():
            renderer.close()
        self._renderers.clear()
        self._temporary.cleanup()

Xet Storage Details

Size:: 27.8 kB
Xet hash:: f0f2f08fdd72d586a0ff3829f6440bb565cd449ced1bedf33cdcdb30e531d9cf

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import json
	import shutil
	import tempfile
	import xml.etree.ElementTree as ET
	from pathlib import Path

	import mujoco
	import numpy as np
	from PIL import Image
	from scipy.spatial import ConvexHull
	from scipy.spatial.transform import Rotation

	from ..geometry import MUJOCO_CAMERA_FROM_OPENCV, quat_wxyz


	class PhysicsScene:
	def __init__(self, folder: Path, asset: Path, reference: Path \| None = None):
	self.folder = Path(folder).resolve()
	self.asset = Path(asset).resolve()
	self._temporary = tempfile.TemporaryDirectory(prefix="real2sim-physics-")
	self._renderers = {}
	with np.load(self.folder / "sim_trajectory.npz") as trajectory:
	self.reference_qpos = trajectory["qpos"].copy()
	self.camera = trajectory["camera"].copy()
	self.timestamps = trajectory["timestamps"].copy()
	self.frame_indices = trajectory["frame_indices"].copy()
	if reference is not None:
	with np.load(reference) as fitted:
	if fitted["qpos"].shape != self.reference_qpos.shape:
	raise ValueError("Fitted hand trajectory must match the reference")
	if "timestamps" in fitted and not np.allclose(
	fitted["timestamps"], self.timestamps
	):
	raise ValueError("Fitted hand timestamps must match the reference")
	self.reference_qpos = fitted["qpos"].copy()
	with np.load(self.folder / "real_masks.npz") as target:
	names = target["names"].tolist()
	self.real_masks = target["masks"][
	[names.index(n) for n in ("left", "right", "object")]
	].copy()
	with np.load(self.asset) as geometry:
	vertices = np.asarray(geometry["world_vertices"], dtype=float)
	self.object_faces = np.asarray(geometry["faces"], dtype=np.int32)
	self.object_colors = np.asarray(geometry["colors"], dtype=np.uint8)
	self.source_frame_index = int(geometry["source_frame_index"])
	if (
	vertices.ndim != 2
	or vertices.shape[1] != 3
	or not np.isfinite(vertices).all()
	):
	raise ValueError("Asset must contain finite world vertices")
	if self.object_faces.ndim != 2 or self.object_faces.shape[1] != 3:
	raise ValueError("Asset must contain triangular faces")
	if self.object_colors.shape != vertices.shape:
	raise ValueError("Asset colors must match vertices")
	if len(self.timestamps) < 2 or np.any(np.diff(self.timestamps) <= 0):
	raise ValueError("Trajectory timestamps must increase")
	self.object_center = (vertices.min(axis=0) + vertices.max(axis=0)) / 2
	self.local_vertices = vertices - self.object_center
	self._prepare_table()
	self.model = self._compile()
	self.data = mujoco.MjData(self.model)
	self.object_joint = self.model.joint("object_free").id
	self.object_qpos_address = int(self.model.jnt_qposadr[self.object_joint])
	self.object_dof_address = int(self.model.jnt_dofadr[self.object_joint])
	self.object_geom = self.model.geom("object_collision").id
	self.object_visual_geom = self.model.geom("object_visual").id
	self.camera_mocap = int(
	self.model.body_mocapid[self.model.body("ego_mount").id]
	)
	self.hand_geoms = []
	self.hand_joints = []
	self.hand_visible = []
	self.hand_mocap = []
	self.hand_equalities = []
	self.wrist_addresses = []
	self.ball_joints = []
	for side in ("left", "right"):
	joints = np.array(
	[
	j
	for j in range(self.model.njnt)
	if self.model.joint(j).name.startswith(side + "_")
	]
	)
	geoms = np.array(
	[
	g
	for g in range(self.model.ngeom)
	if self.model.body(int(self.model.geom_bodyid[g])).name.startswith(
	side + "_"
	)
	]
	)
	root = next(
	j
	for j in joints
	if self.model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE
	)
	address = int(self.model.jnt_qposadr[root])
	self.hand_joints.append(joints)
	self.hand_geoms.append(geoms)
	self.wrist_addresses.append(address)
	self.hand_visible.append(self.reference_qpos[:, address + 2] > -9)
	self.hand_mocap.append(
	int(self.model.body_mocapid[self.model.body(side + "_target").id])
	)
	self.hand_equalities.append(self.model.equality(side + "_weld").id)
	self.ball_joints.extend(
	j
	for j in joints
	if self.model.jnt_type[j] == mujoco.mjtJoint.mjJNT_BALL
	)
	self.ball_joints = np.asarray(self.ball_joints, dtype=int)
	self.actor_names = np.array(["left", "right", "object"])
	self.geom_actor_ids = np.full(self.model.ngeom, -1, dtype=np.int8)
	for actor, geoms in enumerate(self.hand_geoms):
	self.geom_actor_ids[geoms] = actor
	self.geom_actor_ids[[self.object_geom, self.object_visual_geom]] = 2
	self.ball_addresses = self.model.jnt_qposadr[self.ball_joints]
	self._finger_targets = np.stack(
	[
	Rotation.from_quat(
	self.reference_qpos[:, a : a + 4][:, [1, 2, 3, 0]]
	).as_rotvec()
	for a in self.ball_addresses
	],
	axis=1,
	).reshape(len(self.timestamps), -1)
	self.parameter_mean = np.r_[np.zeros(24), 0.8]
	self.parameter_std = np.r_[
	np.full(3, 0.015), np.full(3, 0.12), np.full(18, 0.015), 0.2
	]
	self.parameter_bounds = np.column_stack(
	(
	np.r_[np.full(3, -0.06), np.full(3, -0.5), np.full(18, -0.08), 0.15],
	np.r_[np.full(3, 0.06), np.full(3, 0.5), np.full(18, 0.08), 1.8],
	)
	)
	self.bounds = self.parameter_bounds
	self.parameter_names = [
	"object_" + n for n in ("x", "y", "z", "rx", "ry", "rz")
	]
	self.parameter_names += [
	f"{side}_{key}_{axis}"
	for side in ("left", "right")
	for key in range(3)
	for axis in ("x", "y", "z")
	]
	self.parameter_names += ["friction"]
	self.reset(self.parameter_mean)

	def _prepare_table(self):
	scene = json.loads((self.folder / "scene.json").read_text())
	if scene.get("geometry_source") != "flat_table":
	raise ValueError("Scene geometry_source must be flat_table")
	height = scene.get("table_height_m")
	if (
	isinstance(height, bool)
	or not isinstance(height, (int, float))
	or not np.isfinite(height)
	):
	raise ValueError("Scene table_height_m must be a finite number")
	source = scene.get("table_source")
	if not isinstance(source, str) or not source.strip():
	raise ValueError("Scene table_source must identify the height estimate")
	self.table_height = float(height)
	self.table_diagnostics = {
	"height_m": self.table_height,
	"source": source,
	"geometry_source": "flat_table",
	}

	def _compile(self):
	spec = mujoco.MjSpec.from_file(str(self.folder / "scene.xml"))
	spec.option.timestep = 1 / 600
	spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
	spec.option.iterations = 80
	spec.option.gravity = [0, 0, -9.81]
	spec.visual.global_.offwidth = 1920
	spec.visual.global_.offheight = 1080
	spec.delete(spec.geom("floor"))
	for material in list(spec.materials):
	if material.name == "grid" and not any(
	geom.material == material.name for geom in spec.geoms
	):
	spec.delete(material)
	for texture in list(spec.textures):
	if texture.name == "grid" and not any(
	texture.name in list(material.textures) for material in spec.materials
	):
	spec.delete(texture)
	spec.visual.headlight.ambient = [0.35, 0.35, 0.35]
	spec.visual.headlight.diffuse = [0.55, 0.55, 0.55]
	spec.visual.headlight.specular = [0.05, 0.05, 0.05]
	spec.visual.quality.shadowsize = 2048
	for light in list(spec.lights):
	spec.delete(light)
	spec.worldbody.add_light(
	name="key_light",
	type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL,
	pos=self.object_center + np.array([0.5, -0.5, 2.0]),
	dir=[-0.25, 0.25, -1],
	ambient=[0.1, 0.1, 0.1],
	diffuse=[0.6, 0.6, 0.6],
	specular=[0.05, 0.05, 0.05],
	castshadow=True,
	)
	spec.add_texture(
	name="background",
	type=mujoco.mjtTexture.mjTEXTURE_SKYBOX,
	builtin=mujoco.mjtBuiltin.mjBUILTIN_FLAT,
	rgb1=[0.88, 0.89, 0.91],
	rgb2=[0.88, 0.89, 0.91],
	width=128,
	height=768,
	)
	spec.add_material(
	name="table_material",
	rgba=[0.78, 0.74, 0.65, 1],
	specular=0,
	shininess=0,
	reflectance=0,
	)
	spec.worldbody.add_geom(
	name="table",
	type=mujoco.mjtGeom.mjGEOM_PLANE,
	pos=[0, 0, self.table_height],
	size=[2, 2, 0.1],
	contype=1,
	conaffinity=6,
	material="table_material",
	friction=[0.8, 0.01, 0.001],
	)
	width = int(np.ceil(np.sqrt(len(self.object_faces))))
	height = int(np.ceil(len(self.object_faces) / width))
	pixels = np.zeros((height * width, 3), dtype=np.uint8)
	pixels[: len(self.object_faces)] = (
	self.object_colors[self.object_faces].mean(axis=1).astype(np.uint8)
	)
	self.object_texture = Path(self._temporary.name) / "object_colors.png"
	Image.fromarray(pixels.reshape(height, width, 3)).save(self.object_texture)
	spec.add_texture(
	name="object_colors",
	type=mujoco.mjtTexture.mjTEXTURE_2D,
	file=str(self.object_texture),
	)
	material = spec.add_material(
	name="object_material", texuniform=False, specular=0
	)
	material.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "object_colors"
	atlas = np.arange(len(self.object_faces))
	self.object_uv = np.column_stack(
	((atlas % width + 0.5) / width, (atlas // width + 0.5) / height)
	)
	spec.add_mesh(
	name="object_mesh",
	uservert=self.local_vertices.ravel(),
	userface=self.object_faces.ravel(),
	usertexcoord=self.object_uv.ravel(),
	userfacetexcoord=np.repeat(atlas, 3),
	)
	hull = ConvexHull(self.local_vertices)
	used, faces = np.unique(hull.simplices, return_inverse=True)
	spec.add_mesh(
	name="object_hull",
	uservert=self.local_vertices[used].ravel(),
	userface=faces.ravel(),
	)
	body = spec.worldbody.add_body(name="object", pos=self.object_center)
	body.add_joint(name="object_free", type=mujoco.mjtJoint.mjJNT_FREE)
	body.add_geom(
	name="object_visual",
	type=mujoco.mjtGeom.mjGEOM_MESH,
	meshname="object_mesh",
	material="object_material",
	contype=0,
	conaffinity=0,
	mass=0,
	)
	body.add_geom(
	name="object_collision",
	type=mujoco.mjtGeom.mjGEOM_MESH,
	meshname="object_hull",
	group=3,
	rgba=[0, 0, 0, 0],
	mass=0.1,
	contype=4,
	conaffinity=3,
	condim=4,
	priority=1,
	friction=[0.8, 0.01, 0.001],
	solref=[0.01, 1],
	)
	for geom in spec.geoms:
	if geom.parent.name.startswith(("left_", "right_")):
	geom.contype = 2
	geom.conaffinity = 5
	geom.friction = [0.8, 0.01, 0.001]
	for side in ("left", "right"):
	spec.worldbody.add_body(name=side + "_target", mocap=True)
	spec.add_equality(
	name=side + "_weld",
	type=mujoco.mjtEq.mjEQ_WELD,
	name1=side + "_target",
	name2=side + "_Wrist",
	objtype=mujoco.mjtObj.mjOBJ_BODY,
	data=[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1],
	solref=[0.015, 1],
	)
	for joint in spec.joints:
	if joint.type != mujoco.mjtJoint.mjJNT_BALL:
	continue
	joint.damping = [0.015, 0, 0]
	joint.armature = 0.0001
	for axis in range(3):
	gear = np.zeros(6)
	gear[axis] = 1
	spec.add_actuator(
	name=f"{joint.name}_{axis}",
	trntype=mujoco.mjtTrn.mjTRN_JOINT,
	target=joint.name,
	gear=gear,
	gaintype=mujoco.mjtGain.mjGAIN_FIXED,
	gainprm=[0.8] + [0] * 9,
	biastype=mujoco.mjtBias.mjBIAS_AFFINE,
	biasprm=[0, -0.8, -0.025] + [0] * 7,
	ctrllimited=True,
	ctrlrange=[-np.pi, np.pi],
	forcelimited=True,
	forcerange=[-2, 2],
	)
	self.spec = spec
	return spec.compile()

	def _parameters(self, parameters):
	values = np.asarray(parameters, dtype=float)
	if values.shape != (25,) or not np.isfinite(values).all():
	raise ValueError("Expected 25 finite parameters")
	if np.any(values < self.parameter_bounds[:, 0]) or np.any(
	values > self.parameter_bounds[:, 1]
	):
	raise ValueError("Parameters exceed bounds")
	return values.copy()

	def reset(self, parameters):
	self.parameters = self._parameters(parameters)
	mujoco.mj_resetData(self.model, self.data)
	self.data.qpos[: self.reference_qpos.shape[1]] = self.reference_qpos[0]
	rotation = Rotation.from_rotvec(self.parameters[3:6])
	quaternion = rotation.as_quat()[[3, 0, 1, 2]]
	position = self.object_center + self.parameters[:3]
	minimum_z = float((self.local_vertices @ rotation.as_matrix().T)[:, 2].min())
	self.initial_table_adjustment = max(
	0.0, self.table_height + 0.001 - (position[2] + minimum_z)
	)
	position[2] += self.initial_table_adjustment
	address = self.object_qpos_address
	self.data.qpos[address : address + 7] = np.r_[position, quaternion]
	self.model.geom_friction[self.object_geom, 0] = self.parameters[-1]
	self._active = np.zeros(2, dtype=bool)
	self._contact_start = np.full(2, np.inf)
	self._activation_frames = [[], []]
	self._collision_enabled = np.zeros(2, dtype=bool)
	self._contact_steps = 0
	self._contact_pairs = 0
	self._contact_peak_force = 0.0
	self._contact_impulse = 0.0
	self._contact_by_hand = np.zeros(2, dtype=int)
	self._initial_object_position = position.copy()
	self._maximum_object_z = position[2]
	self._unstable = False
	for geoms in self.hand_geoms:
	self.model.geom_contype[geoms] = 0
	self.model.geom_conaffinity[geoms] = 0
	self._drive(0.0)
	mujoco.mj_forward(self.model, self.data)

	def _drive(self, time):
	high = min(
	int(np.searchsorted(self.timestamps, time, side="right")),
	len(self.timestamps) - 1,
	)
	low = max(0, high - 1)
	fraction = np.clip(
	(time - self.timestamps[low])
	/ (self.timestamps[high] - self.timestamps[low]),
	0,
	1,
	)
	nearest = min(
	int(np.searchsorted(self.timestamps, time + 1e-9, side="right")) - 1,
	len(self.timestamps) - 1,
	)
	nearest = max(0, nearest)
	control = (1 - fraction) * self._finger_targets[
	low
	] + fraction * self._finger_targets[high]
	self.data.ctrl[:] = control
	corrections = self.parameters[6:24].reshape(2, 3, 3)
	key_times = np.linspace(self.timestamps[0], self.timestamps[-1], 3)
	for side, address in enumerate(self.wrist_addresses):
	visible = bool(self.hand_visible[side][nearest])
	mocap = self.hand_mocap[side]
	if visible:
	correction = np.array(
	[
	np.interp(time, key_times, corrections[side, :, axis])
	for axis in range(3)
	]
	)
	a, b = self.reference_qpos[[low, high], address : address + 7].copy()
	if not self.hand_visible[side][low]:
	a = b.copy()
	if not self.hand_visible[side][high]:
	b = a.copy()
	position = (1 - fraction) * a[:3] + fraction * b[:3] + correction
	qb = b[3:] * (1 if np.dot(a[3:], b[3:]) >= 0 else -1)
	quaternion = (1 - fraction) * a[3:] + fraction * qb
	quaternion /= np.linalg.norm(quaternion)
	if not self._active[side]:
	for joint in self.hand_joints[side]:
	start = int(self.model.jnt_qposadr[joint])
	size = (
	7
	if self.model.jnt_type[joint] == mujoco.mjtJoint.mjJNT_FREE
	else 4
	)
	self.data.qpos[start : start + size] = self.reference_qpos[
	nearest, start : start + size
	]
	velocity = int(self.model.jnt_dofadr[joint])
	count = 6 if size == 7 else 3
	self.data.qvel[velocity : velocity + count] = 0
	self.data.qpos[address : address + 7] = np.r_[position, quaternion]
	self._contact_start[side] = time + 0.1
	self._activation_frames[side].append(nearest)
	enabled = time >= self._contact_start[side]
	else:
	position = np.array([5.0 + side, 5.0, self.table_height + 1])
	quaternion = np.array([1.0, 0, 0, 0])
	enabled = False
	if self._active[side] or time == 0:
	self.data.qpos[address : address + 7] = np.r_[position, quaternion]
	root = int(self.model.jnt_dofadr[self.hand_joints[side][0]])
	self.data.qvel[root : root + 6] = 0
	self.data.mocap_pos[mocap] = position
	self.data.mocap_quat[mocap] = quaternion
	if enabled != self._collision_enabled[side]:
	self.model.geom_contype[self.hand_geoms[side]] = 2 if enabled else 0
	self.model.geom_conaffinity[self.hand_geoms[side]] = 5 if enabled else 0
	self._collision_enabled[side] = enabled
	self._active[side] = visible

	def _contacts(self):
	force = np.zeros(6)
	touched = False
	for index in range(self.data.ncon):
	contact = self.data.contact[index]
	if contact.geom1 == self.object_geom:
	other = contact.geom2
	elif contact.geom2 == self.object_geom:
	other = contact.geom1
	else:
	continue
	for side, geoms in enumerate(self.hand_geoms):
	if other not in geoms:
	continue
	mujoco.mj_contactForce(self.model, self.data, index, force)
	normal = max(0.0, float(force[0]))
	if normal <= 1e-8:
	continue
	touched = True
	self._contact_pairs += 1
	self._contact_by_hand[side] += 1
	self._contact_peak_force = max(self._contact_peak_force, normal)
	self._contact_impulse += normal * self.model.opt.timestep
	self._contact_steps += int(touched)
	self._maximum_object_z = max(
	self._maximum_object_z, float(self.data.qpos[self.object_qpos_address + 2])
	)

	def _renderer(self, size):
	size = tuple(size)
	if size not in self._renderers:
	self._renderers[size] = mujoco.Renderer(
	self.model, width=size[0], height=size[1]
	)
	return self._renderers[size]

	def rollout(self, parameters, frame_indices, render_size=(256, 144), rgb=False):
	indices = np.asarray(frame_indices, dtype=int)
	if (
	indices.ndim != 1
	or not len(indices)
	or np.any(np.diff(indices) <= 0)
	or indices[0] < 0
	or indices[-1] >= len(self.timestamps)
	):
	raise ValueError("Frame indices must increase within the saved trajectory")
	self.reset(parameters)
	renderer = self._renderer(render_size)
	masks, images, qpos, positions, full_qpos, counts, forces, geom_ids = (
	[],
	[],
	[],
	[],
	[],
	[],
	[],
	[],
	)
	for frame in indices:
	target_time = self.timestamps[frame] - self.timestamps[0]
	while self.data.time + self.model.opt.timestep / 2 < target_time:
	self._drive(self.data.time + self.timestamps[0])
	mujoco.mj_step(self.model, self.data)
	self._contacts()
	if (
	not np.isfinite(self.data.qpos).all()
	or np.max(np.abs(self.data.qvel)) > 1e4
	):
	self._unstable = True
	raise FloatingPointError("Physics rollout became unstable")
	self._drive(self.data.time + self.timestamps[0])
	self.data.mocap_pos[self.camera_mocap] = self.camera[frame, :3, 3]
	self.data.mocap_quat[self.camera_mocap] = quat_wxyz(
	self.camera[frame, :3, :3] @ MUJOCO_CAMERA_FROM_OPENCV
	)
	mujoco.mj_forward(self.model, self.data)
	renderer.enable_segmentation_rendering()
	renderer.update_scene(self.data, camera="ego")
	segmentation = renderer.render()
	is_geom = segmentation[..., 1] == mujoco.mjtObj.mjOBJ_GEOM
	ids = segmentation[..., 0]
	geom_ids.append(np.where(is_geom, ids, -1).astype(np.int32))
	masks.append(
	np.stack(
	[
	is_geom & np.isin(ids, geoms)
	for geoms in (*self.hand_geoms, [self.object_visual_geom])
	]
	)
	)
	if rgb:
	renderer.disable_segmentation_rendering()
	renderer.update_scene(self.data, camera="ego")
	images.append(renderer.render().copy())
	address = self.object_qpos_address
	qpos.append(self.data.qpos[address : address + 7].copy())
	positions.append(self.data.qpos[address : address + 3].copy())
	full_qpos.append(self.data.qpos.copy())
	counts.append(self._contact_pairs)
	forces.append(self._contact_peak_force)
	diagnostics = {
	"simulation_mode": "contact_physics",
	"hand_object_contact_steps": self._contact_steps,
	"hand_object_contact_pairs": self._contact_pairs,
	"hand_object_contacts_by_hand": self._contact_by_hand.tolist(),
	"peak_hand_object_normal_force_n": self._contact_peak_force,
	"hand_object_normal_impulse_ns": self._contact_impulse,
	"maximum_lift_m": float(
	self._maximum_object_z - self._initial_object_position[2]
	),
	"final_displacement_m": float(
	np.linalg.norm(positions[-1] - self._initial_object_position)
	),
	"initial_table_adjustment_m": float(self.initial_table_adjustment),
	"hand_activation_frames": self._activation_frames,
	"activation_collision_delay_s": 0.1,
	"table": self.table_diagnostics,
	"unstable": self._unstable,
	}
	result = {
	"masks": np.stack(masks, axis=1),
	"geom_ids": np.stack(geom_ids),
	"geom_actor_ids": self.geom_actor_ids.copy(),
	"actor_names": self.actor_names.copy(),
	"object_qpos": np.stack(qpos),
	"object_positions": np.stack(positions),
	"qpos": np.stack(full_qpos),
	"frame_indices": indices.copy(),
	"timestamps": self.timestamps[indices].copy(),
	"contact_counts": np.asarray(counts),
	"contact_forces": np.asarray(forces),
	"diagnostics": diagnostics,
	}
	if rgb:
	result["rgb"] = np.stack(images)
	return result

	def save(self, folder):
	folder = Path(folder)
	folder.mkdir(parents=True, exist_ok=True)
	mujoco.mj_saveModel(self.model, str(folder / "scene.mjb"), None)
	shutil.copyfile(self.object_texture, folder / "object_colors.png")
	with (folder / "object_visual.obj").open("w") as mesh:
	np.savetxt(mesh, self.local_vertices, fmt="v %.9g %.9g %.9g")
	np.savetxt(mesh, self.object_uv, fmt="vt %.9g %.9g")
	for index, face in enumerate(self.object_faces, start=1):
	mesh.write(
	"f " + " ".join(f"{vertex + 1}/{index}" for vertex in face) + "\n"
	)
	xml = self.spec.to_xml().replace(str(self.object_texture), "object_colors.png")
	root = ET.fromstring(xml)
	for mesh in root.findall("asset/mesh"):
	if mesh.get("name") == "object_mesh":
	mesh.attrib.clear()
	mesh.set("name", "object_mesh")
	mesh.set("file", "object_visual.obj")
	for geom in root.findall(".//geom"):
	if geom.get("name") == "object_collision":
	geom.set("friction", f"{self.parameters[-1]:.9g} 0.01 0.001")
	(folder / "scene.xml").write_text(ET.tostring(root, encoding="unicode"))
	np.savez_compressed(
	folder / "controller.npz",
	reference_qpos=self.reference_qpos,
	camera=self.camera,
	timestamps=self.timestamps,
	parameters=self.parameters,
	parameter_mean=self.parameter_mean,
	parameter_std=self.parameter_std,
	parameter_bounds=self.parameter_bounds,
	initial_object_position=self._initial_object_position,
	)
	(folder / "scene_metrics.json").write_text(
	json.dumps(
	{
	"table": self.table_diagnostics,
	"mass_kg": 0.1,
	"source_frame_index": self.source_frame_index,
	"simulation_mode": "contact_physics",
	},
	indent=2,
	)
	)

	def close(self):
	for renderer in self._renderers.values():
	renderer.close()
	self._renderers.clear()
	self._temporary.cleanup()