Buckets:
| 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.ndimage import binary_dilation | |
| from scipy.spatial import ConvexHull | |
| from scipy.spatial.transform import Rotation | |
| from ..geometry import MUJOCO_CAMERA_FROM_OPENCV, quat_wxyz | |
| from ..reconstruction.surface import transform_points, triangulate | |
| 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_surface() | |
| 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_surface(self): | |
| with np.load(self.folder / "view_0000.npz") as source: | |
| view = {key: source[key] for key in source.files} | |
| with np.load(self.folder / "surface.npz") as surface: | |
| transform = surface["T_world_reconstruction"] | |
| excluded = binary_dilation(self.real_masks[:, 0].any(axis=0), iterations=5) | |
| vertices, faces, uv = triangulate(view, exclude=excluded) | |
| self.surface_vertices = transform_points(vertices, transform) | |
| self.surface_faces = faces | |
| self.surface_uv = uv | |
| points = transform_points(view["pts3d"], transform) | |
| valid = view["mask"] & ~excluded & np.isfinite(points).all(axis=-1) | |
| dy, dx = np.gradient(points, axis=(0, 1)) | |
| normals = np.cross(dx, dy) | |
| norm = np.linalg.norm(normals, axis=-1) | |
| horizontal = np.abs(normals[..., 2]) / np.maximum(norm, 1e-12) > 0.85 | |
| distance = np.linalg.norm(points[..., :2] - self.object_center[:2], axis=-1) | |
| selected = ( | |
| valid | |
| & horizontal | |
| & (distance < 0.4) | |
| & (points[..., 2] < self.object_center[2] + 0.02) | |
| ) | |
| heights = points[..., 2][selected] | |
| if len(heights) < 30: | |
| raise ValueError("Insufficient mapped table points near the object") | |
| lower, upper = np.quantile(heights, [0.05, 0.95]) | |
| edges = np.arange(lower, upper + 0.005, 0.005) | |
| if len(edges) < 2: | |
| edges = np.array([lower - 0.0025, lower + 0.0025]) | |
| histogram, edges = np.histogram(heights, bins=edges) | |
| peak = (edges[np.argmax(histogram)] + edges[np.argmax(histogram) + 1]) / 2 | |
| inliers = np.abs(heights - peak) < 0.01 | |
| self.table_height = float(np.median(heights[inliers])) | |
| self.table_diagnostics = { | |
| "height_m": self.table_height, | |
| "mapped_points": len(heights), | |
| "inlier_points": int(inliers.sum()), | |
| "inlier_height_mad_m": float( | |
| np.median(np.abs(heights[inliers] - self.table_height)) | |
| ), | |
| } | |
| def _compile(self): | |
| spec = mujoco.MjSpec.from_file(str(self.folder / "scene.xml")) | |
| spec.textures[0].file = str(self.folder / "view_0000.png") | |
| 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("scene_surface")) | |
| spec.delete(spec.mesh("scene_surface")) | |
| spec.delete(spec.geom("floor")) | |
| spec.add_mesh( | |
| name="scene_surface", | |
| uservert=self.surface_vertices.ravel(), | |
| userface=self.surface_faces.ravel(), | |
| usertexcoord=self.surface_uv.ravel(), | |
| userfacetexcoord=self.surface_faces.ravel(), | |
| ) | |
| spec.worldbody.add_geom( | |
| name="scene_surface", | |
| type=mujoco.mjtGeom.mjGEOM_MESH, | |
| meshname="scene_surface", | |
| material="scene_rgb", | |
| contype=0, | |
| conaffinity=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, | |
| rgba=[0.68, 0.65, 0.58, 1], | |
| 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.folder / "view_0000.png", folder / "view_0000.png") | |
| 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.folder / "view_0000.png"), "view_0000.png") | |
| .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:
- 28.4 kB
- Xet hash:
- eea31638d8fd7ee348dc15411dd47655003d6c33e2224e3a6e26b3caf7ba56f7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.