| """Authored contact-based rigid-block task on the pinned native six-axis reference. |
| |
| The policy is a development baseline, not a learned skill. Observations are |
| privileged simulator state. Only reset initializes the freely moving block; |
| grasping, transport and release thereafter require native contact dynamics. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| from hashlib import sha256 |
| import json |
| import math |
| from pathlib import Path |
| import time |
|
|
| from . import reference_sim as sim |
|
|
|
|
| PROFILE = 'ur5e-2f85-rigid-block-task-v0.1' |
| SCHEMA = 'dvidia.reference-task-episode/v1' |
| SCOPE = ('Authored native UR5e/Robotiq rigid-block manipulation, privileged state and synthetic commands. ' |
| 'No shoe model, learned task policy, camera sensing, calibrated contact or physical qualification.') |
| CARRY_CONTACT_LOSS_TOLERANCE_S = .04 |
| GRIPPER_JOINTS = tuple('gripper/' + side + '_' + link + '_joint' |
| for side in ('right', 'left') |
| for link in ('driver', 'coupler', 'spring_link', 'follower')) |
|
|
|
|
| def _number(value, name, minimum, maximum): |
| if isinstance(value, bool) or not isinstance(value, (int, float)): |
| raise ValueError(f'{name} must be finite numeric data') |
| try: |
| value = float(value) |
| except (OverflowError, ValueError, TypeError) as exc: |
| raise ValueError(f'{name} must be finite numeric data') from exc |
| if not math.isfinite(value) or not minimum <= value <= maximum: |
| raise ValueError(f'{name} must lie in [{minimum}, {maximum}]') |
| return value |
|
|
|
|
| def _vector(value, length, name, minimum=-10., maximum=10.): |
| if not isinstance(value, (list, tuple)) or len(value) != length: |
| raise ValueError(f'{name} requires {length} values') |
| return tuple(_number(item, name, minimum, maximum) for item in value) |
|
|
|
|
| @dataclass(frozen=True) |
| class ReferenceTaskConfig: |
| timestep: float = .001 |
| control_dt: float = .02 |
| timeout_s: float = 20. |
| settle_s: float = .5 |
| object_position_m: tuple = (-.134, .49, .302) |
| target_position_m: tuple = (.08, .49, .30) |
| object_half_size_m: tuple = (.02, .02, .02) |
| mass_kg: float = .04 |
| friction: float = .8 |
| table_top_m: float = .28 |
| table_center_xy_m: tuple = (0., .55) |
| table_half_size_xy_m: tuple = (.32, .26) |
| table_thickness_m: float = .04 |
| goal_position_tolerance_m: float = .025 |
| goal_dwell_s: float = .25 |
| min_lift_m: float = .06 |
| min_lift_dwell_s: float = .10 |
| retreat_clearance_m: float = .10 |
| joint_target_rate_rad_s: float = 1.5 |
| jaw_target_rate_0_255_s: float = 500. |
| tool_orientation_wxyz: tuple = (0., math.sqrt(.5), math.sqrt(.5), 0.) |
|
|
| def __post_init__(self): |
| bounds = {'timestep': (.0005, .002), 'control_dt': (.02, .02), 'timeout_s': (.1, 30.), |
| 'settle_s': (.1, 2.), 'mass_kg': (.01, .15), 'friction': (.02, 2.), |
| 'table_top_m': (.24, .32), 'table_thickness_m': (.02, .08), |
| 'goal_position_tolerance_m': (.01, .03), 'goal_dwell_s': (.10, .6), |
| 'min_lift_m': (.04, .10), 'min_lift_dwell_s': (.05, .3), |
| 'retreat_clearance_m': (.08, .15), 'joint_target_rate_rad_s': (.1, 3.), |
| 'jaw_target_rate_0_255_s': (50., 1000.)} |
| for name, (low, high) in bounds.items(): |
| object.__setattr__(self, name, _number(getattr(self, name), name, low, high)) |
| sim._timestep(self.timestep) |
| for numerator, denominator in ((self.control_dt, self.timestep), (self.timeout_s, self.control_dt), |
| (self.settle_s, self.timestep)): |
| ratio = numerator / denominator |
| if not math.isclose(ratio, round(ratio), abs_tol=1e-9, rel_tol=0): |
| raise ValueError('Task durations must contain an integer number of native/control ticks') |
| for name, length, low, high in (('object_position_m', 3, -2, 2), ('target_position_m', 3, -2, 2), |
| ('object_half_size_m', 3, .015, .025), ('table_center_xy_m', 2, -1, 1), |
| ('table_half_size_xy_m', 2, .1, .5), ('tool_orientation_wxyz', 4, -1, 1)): |
| object.__setattr__(self, name, _vector(getattr(self, name), length, name, low, high)) |
| if not math.isclose(sum(x*x for x in self.tool_orientation_wxyz), 1., abs_tol=1e-8, rel_tol=0): |
| raise ValueError('Tool goal quaternion must be normalized wxyz') |
| for name in ('object_position_m', 'target_position_m'): |
| point = getattr(self, name) |
| if not -.25 <= point[0] <= .25 or not .35 <= point[1] <= .70: |
| raise ValueError('Task placement lies outside the declared reference reach envelope') |
| for axis in range(2): |
| if abs(point[axis]-self.table_center_xy_m[axis]) + self.object_half_size_m[axis] > self.table_half_size_xy_m[axis]-.01: |
| raise ValueError('Task placement must be supported within the table edge') |
| expected = self.table_top_m + self.object_half_size_m[2] + (.002 if name == 'object_position_m' else 0.) |
| if not math.isclose(point[2], expected, abs_tol=1e-8, rel_tol=0): |
| raise ValueError('Start/target height must match the declared support plane and block size') |
| if math.dist(self.object_position_m[:2], self.target_position_m[:2]) < .075: |
| raise ValueError('Task must require at least 75 mm of block transfer') |
|
|
|
|
| def task_contract(): |
| return {'profile': PROFILE, 'scope': SCOPE, 'physical_robot_ready': False, |
| 'arm_joint_order': list(sim.ARM_JOINTS), 'gripper_linkage_order': list(GRIPPER_JOINTS), |
| 'actions': {'joint_targets_rad': 'six absolute joint targets, rad; ramped then gravity feedforward', |
| 'gripper_0_255': 'one native tendon servo target; 0 open / 255 close, not measured gap or force'}, |
| 'observation_origin': 'privileged_native_simulator_state', 'native_sensor_definitions': 0, |
| 'observations': {'joint_position_rad': [6], 'joint_velocity_rad_s': [6], 'linkage_position_rad': [8], |
| 'tcp_position_m': [3], 'tcp_quaternion_wxyz': [4], 'object_position_m': [3], |
| 'object_quaternion_wxyz': [4], 'grasp_contact_counts': 'left/right native block contacts', |
| 'pad_normal_forces_n': 'native mj_contactForce, not deployable tactile sensing'}, |
| 'contacts': 'native collision proxies; object priority 1 overrides pad/table friction floors', |
| 'object': 'one free rigid box; no weld, grasp attachment or runtime object pose writes', |
| 'success': 'valid physics, retained bilateral table-free lift and supported destination carry completion, then released/settled destination dwell and tool retreat', |
| 'carry_retention': {'bilateral_contact_required': 'after sustained lift until supported at destination while held', |
| 'contact_force_threshold_n': .01, |
| 'consecutive_loss_tolerance_seconds': CARRY_CONTACT_LOSS_TOLERANCE_S, |
| 'loss_at_or_above_tolerance': 'permanent failure, independent of teacher phase'}, |
| 'fixed_base_collision_proxy': False} |
|
|
|
|
| def build_task_model(assets, config=None): |
| config = ReferenceTaskConfig() if config is None else config |
| if not isinstance(config, ReferenceTaskConfig): |
| raise TypeError('config must be ReferenceTaskConfig') |
| _, blobs = sim._asset_payload(assets) |
| mj, _ = sim._native_imports() |
| arm = sim._spec(blobs, 'universal_robots_ur5e', 'ur5e.xml') |
| jaw = sim._spec(blobs, 'robotiq_2f85', '2f85.xml') |
| arm.option.cone, arm.option.impratio = jaw.option.cone, jaw.option.impratio |
| arm.option.timestep = jaw.option.timestep = config.timestep |
| arm.option.integrator = jaw.option.integrator = mj.mjtIntegrator.mjINT_IMPLICITFAST |
| arm.attach(jaw, prefix='gripper/', site='attachment_site') |
| arm.worldbody.add_geom(name='task_floor', type=mj.mjtGeom.mjGEOM_PLANE, size=[2, 2, .1], rgba=[.9, .9, .88, 1]) |
| arm.worldbody.add_geom(name='task_table', type=mj.mjtGeom.mjGEOM_BOX, |
| pos=[*config.table_center_xy_m, config.table_top_m-config.table_thickness_m/2], |
| size=[*config.table_half_size_xy_m, config.table_thickness_m/2], rgba=[.73, .70, .64, 1]) |
| block = arm.worldbody.add_body(name='task_block', pos=config.object_position_m) |
| block.add_freejoint(name='task_block_free') |
| block.add_geom(name='task_block_geom', type=mj.mjtGeom.mjGEOM_BOX, size=config.object_half_size_m, |
| mass=config.mass_kg, friction=[config.friction, .005, .0001], priority=1, |
| condim=4, rgba=[.75, .30, .16, 1]) |
| marker = arm.worldbody.add_body(name='task_target_frame', |
| pos=[*config.target_position_m[:2], config.table_top_m+.0005]) |
| marker.add_geom(name='task_target', type=mj.mjtGeom.mjGEOM_BOX, |
| size=[config.object_half_size_m[0]+.015, config.object_half_size_m[1]+.015, .0005], |
| rgba=[.17, .47, .35, .3], contype=0, conaffinity=0) |
| model = arm.compile() |
| if (model.nq, model.nv, model.nu, model.neq) != (21, 20, 7, 3): |
| raise ValueError('Task native mechanical contract changed') |
| return model |
|
|
|
|
| class ReferenceTaskEnv: |
| def __init__(self, assets, config=None): |
| self.assets = Path(assets) |
| self.config = ReferenceTaskConfig() if config is None else config |
| if not isinstance(self.config, ReferenceTaskConfig): |
| raise TypeError('config must be ReferenceTaskConfig') |
| self.ready = False |
|
|
| def reset(self, seed=0): |
| if type(seed) is not int or not 0 <= seed <= 999999: |
| raise ValueError('seed must be an integer in [0, 999999]') |
| mj, np = sim._native_imports() |
| self.ready = False |
| self.seed, self.warnings = seed, [] |
| with sim.native_warnings(self.warnings): |
| self.model = build_task_model(self.assets, self.config) |
| self.data = mj.MjData(self.model) |
| self.arm_ids = np.asarray([mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_JOINT, x) for x in sim.ARM_JOINTS]) |
| self.arm_qpos = self.model.jnt_qposadr[self.arm_ids] |
| self.arm_dof = self.model.jnt_dofadr[self.arm_ids] |
| self.linkage_ids = np.asarray([mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_JOINT, x) for x in GRIPPER_JOINTS]) |
| self.linkage_qpos = self.model.jnt_qposadr[self.linkage_ids] |
| self.tcp = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_SITE, 'gripper/pinch') |
| self.block = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_BODY, 'task_block') |
| self.block_geom = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_GEOM, 'task_block_geom') |
| self.table_geom = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_GEOM, 'task_table') |
| self.floor_geom = mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_GEOM, 'task_floor') |
| self.pad_geoms = {side: [mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_GEOM, f'gripper/{side}_pad{i}') |
| for i in (1, 2)] for side in ('left', 'right')} |
| self.robot_geoms = set(range(self.model.ngeom)) - { |
| self.block_geom, self.table_geom, |
| self.floor_geom, |
| mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_GEOM, 'task_target')} |
| arm_bodies = {mj.mj_name2id(self.model, mj.mjtObj.mjOBJ_BODY, name) for name in sim.ARM_BODIES} |
| self.arm_geoms = {i for i in self.robot_geoms if int(self.model.geom_bodyid[i]) in arm_bodies} |
| |
| |
| self.data.qpos[self.arm_qpos] = [-1.5708, -1.5708, 1.5708, -1.5708, -1.5708, 0.] |
| self.data.qpos[self.linkage_qpos] = 0. |
| self.command_targets = self.data.qpos[self.arm_qpos].copy() |
| self.command_gripper = 0. |
| mj.mj_forward(self.model, self.data) |
| for _ in range(round(self.config.settle_s/self.config.timestep)): |
| self._controls() |
| mj.mj_step(self.model, self.data) |
| mj.mj_forward(self.model, self.data) |
| if not self._finite() or self.warnings or any(x.number for x in self.data.warning): |
| raise RuntimeError('Task settling produced invalid native physics') |
| self.initial_object_position = self.data.xpos[self.block].copy() |
| self.elapsed, self.steps, self.dwell, self.lift_dwell = 0., 0, 0., 0. |
| self.max_lift, self.max_pad_force, self.critical_contact_ticks = 0., 0., 0 |
| self.lift_seen, self.grasp_seen, self.done = False, False, False |
| self.carry_complete, self.retention_lost = False, False |
| self.carry_contact_loss, self.max_carry_contact_loss = 0., 0. |
| self.carry_contact_loss_events = 0 |
| self.first_success_time = None |
| self.reason = 'running' |
| self.ready = True |
| return self.observation(), self.info() |
|
|
| def _finite(self): |
| _, np = sim._native_imports() |
| return all(np.isfinite(value).all() for value in (self.data.qpos, self.data.qvel, self.data.qacc, self.data.xpos)) |
|
|
| def _controls(self): |
| _, np = sim._native_imports() |
| |
| |
| gains = self.model.actuator_gainprm[:6, 0] |
| requested = self.command_targets + self.data.qfrc_bias[self.arm_dof]/gains |
| self.data.ctrl[:6] = np.clip(requested, self.model.actuator_ctrlrange[:6, 0], self.model.actuator_ctrlrange[:6, 1]) |
| self.data.ctrl[6] = self.command_gripper |
|
|
| def _contacts(self): |
| mj, np = sim._native_imports() |
| counts, forces, table, critical = {'left': 0, 'right': 0}, {'left': 0., 'right': 0.}, 0, 0 |
| pads = set(sum(self.pad_geoms.values(), [])) |
| for index, contact in enumerate(self.data.contact[:self.data.ncon]): |
| pair = {int(contact.geom1), int(contact.geom2)} |
| force = np.zeros(6) |
| mj.mj_contactForce(self.model, self.data, index, force) |
| if force[0] <= .01: |
| continue |
| if self.block_geom in pair: |
| if self.table_geom in pair: |
| table += 1 |
| for side in ('left', 'right'): |
| if pair.intersection(self.pad_geoms[side]): |
| counts[side] += 1 |
| forces[side] += float(force[0]) |
| if pair.intersection(self.robot_geoms-pads): |
| critical += 1 |
| elif pair.intersection({self.table_geom, self.floor_geom}) and pair.intersection(self.robot_geoms): |
| critical += 1 |
| elif pair.intersection(self.arm_geoms) and pair.issubset(self.robot_geoms): |
| critical += 1 |
| return counts, forces, table, critical |
|
|
| def observation(self): |
| mj, np = sim._native_imports() |
| if not self._finite(): |
| raise RuntimeError('Cannot record a nonfinite native state') |
| counts, forces, table, critical = self._contacts() |
| tcp_quat, velocity = np.empty(4), np.empty(6) |
| mj.mju_mat2Quat(tcp_quat, self.data.site_xmat[self.tcp]) |
| mj.mj_objectVelocity(self.model, self.data, mj.mjtObj.mjOBJ_BODY, self.block, velocity, 0) |
| separation = float(np.linalg.norm(self.data.geom_xpos[self.pad_geoms['left'][0]]- |
| self.data.geom_xpos[self.pad_geoms['right'][0]])) |
| return {'t_seconds': self.elapsed, 'qpos': self.data.qpos.tolist(), 'qvel': self.data.qvel.tolist(), |
| 'joint_position_rad': self.data.qpos[self.arm_qpos].tolist(), |
| 'joint_velocity_rad_s': self.data.qvel[self.arm_dof].tolist(), |
| 'linkage_position_rad': self.data.qpos[self.linkage_qpos].tolist(), |
| 'tcp_position_m': self.data.site_xpos[self.tcp].tolist(), 'tcp_quaternion_wxyz': tcp_quat.tolist(), |
| 'object_position_m': self.data.xpos[self.block].tolist(), 'object_quaternion_wxyz': self.data.xquat[self.block].tolist(), |
| 'object_linear_velocity_m_s': velocity[3:].tolist(), 'object_angular_velocity_rad_s': velocity[:3].tolist(), |
| 'target_position_m': list(self.config.target_position_m), 'pad_center_separation_m': separation, |
| 'grasp_contact_counts': counts, 'pad_normal_forces_n': forces, 'table_object_contact_count': table, |
| 'unexpected_contact_count': critical, 'contact_count': int(self.data.ncon), |
| 'native_actuator_controls': self.data.ctrl.tolist(), 'actuator_force': self.data.actuator_force.tolist(), |
| 'issued_joint_targets_rad': None, 'issued_gripper_0_255': None, 'command': None, |
| 'observation_origin': 'privileged_native_simulator_state'} |
|
|
| def info(self): |
| _, np = sim._native_imports() |
| valid = self.reason not in ('invalid_state', 'solver_warning') and not self.warnings |
| success = valid and self.carry_complete and not self.retention_lost and self.dwell+1e-12 >= self.config.goal_dwell_s |
| return {'valid': valid, 'success': success, |
| 'reason': ('success' if success else self.reason), 'simulation_time': self.elapsed, |
| 'steps': self.steps, 'lift_seen': self.lift_seen, 'grasp_seen': self.grasp_seen, |
| 'max_object_lift_m': self.max_lift, 'dwell_seconds': self.dwell, 'lift_dwell_seconds': self.lift_dwell, |
| 'max_pad_normal_force_n': self.max_pad_force, 'unexpected_contact_native_ticks': self.critical_contact_ticks, |
| 'carry_complete': self.carry_complete, 'retention_lost': self.retention_lost, |
| 'max_carry_contact_loss_seconds': self.max_carry_contact_loss, |
| 'carry_contact_loss_events': self.carry_contact_loss_events, |
| 'carry_contact_loss_tolerance_seconds': CARRY_CONTACT_LOSS_TOLERANCE_S, |
| 'first_success_time_seconds': self.first_success_time, |
| 'target_distance_m': float(np.linalg.norm(self.data.xpos[self.block]-self.config.target_position_m)) if self._finite() else None, |
| 'warnings': list(self.warnings), 'physical_robot_ready': False} |
|
|
| def _outcome_tick(self, *, both, supported, lift, distance, released, settled, clear, critical, pad_force): |
| """Score native evidence without consulting the authored policy phase.""" |
| dt = self.config.timestep |
| self.max_pad_force = max(self.max_pad_force, pad_force) |
| self.critical_contact_ticks += bool(critical) |
| self.grasp_seen |= both |
| self.max_lift = max(self.max_lift, lift) |
| loaded = both and not supported and lift >= self.config.min_lift_m |
| self.lift_dwell = self.lift_dwell+dt if loaded else 0. |
| self.lift_seen |= self.lift_dwell+1e-12 >= self.config.min_lift_dwell_s |
| if self.lift_seen and not self.carry_complete: |
| |
| |
| self.carry_complete |= bool(both and supported and distance <= self.config.goal_position_tolerance_m) |
| if not self.carry_complete: |
| if both: |
| self.carry_contact_loss = 0. |
| else: |
| before = self.carry_contact_loss |
| self.carry_contact_loss += dt |
| self.max_carry_contact_loss = max(self.max_carry_contact_loss, self.carry_contact_loss) |
| if before+1e-12 < CARRY_CONTACT_LOSS_TOLERANCE_S <= self.carry_contact_loss+1e-12: |
| self.carry_contact_loss_events += 1 |
| self.retention_lost = True |
| qualifies = (self.lift_seen and self.grasp_seen and self.carry_complete and not self.retention_lost |
| and released and supported and settled and clear |
| and distance <= self.config.goal_position_tolerance_m and self.critical_contact_ticks == 0) |
| self.dwell = self.dwell+dt if qualifies else 0. |
| if self.dwell+1e-12 >= self.config.goal_dwell_s and self.first_success_time is None: |
| self.first_success_time = self.elapsed |
|
|
| def pose_action(self, position_m, quaternion_wxyz, gripper_0_255): |
| if not self.ready or self.done: |
| raise RuntimeError('Reset is required before another task action') |
| mj, np = sim._native_imports() |
| position = np.asarray(_vector(position_m, 3, 'pose position'), dtype=float) |
| quaternion = np.asarray(_vector(quaternion_wxyz, 4, 'pose quaternion', -1, 1), dtype=float) |
| if not math.isclose(float(quaternion@quaternion), 1., abs_tol=1e-8, rel_tol=0): |
| raise ValueError('Pose quaternion must be normalized wxyz') |
| jaw = _number(gripper_0_255, 'gripper command', 0, 255) |
| jacp, jacr = np.zeros((3, self.model.nv)), np.zeros((3, self.model.nv)) |
| mj.mj_jacSite(self.model, self.data, jacp, jacr, self.tcp) |
| current, rotation = np.empty(4), np.empty(3) |
| mj.mju_mat2Quat(current, self.data.site_xmat[self.tcp]) |
| mj.mju_subQuat(rotation, quaternion, current) |
| |
| |
| rotation = self.data.site_xmat[self.tcp].reshape(3, 3) @ rotation |
| translation = position-self.data.site_xpos[self.tcp] |
| for error, maximum in ((translation, .018), (rotation, .04)): |
| norm = float(np.linalg.norm(error)) |
| if norm > maximum: |
| error *= maximum/norm |
| jacobian = np.vstack((jacp[:, self.arm_dof], .25*jacr[:, self.arm_dof])) |
| error = np.concatenate((translation, .25*rotation)) |
| delta = jacobian.T @ np.linalg.solve(jacobian@jacobian.T+.012**2*np.eye(6), error) |
| delta = np.clip(delta, -.06, .06) |
| targets = np.clip(self.data.qpos[self.arm_qpos]+delta, |
| self.model.actuator_ctrlrange[:6, 0]+.02, self.model.actuator_ctrlrange[:6, 1]-.02) |
| return {'joint_targets_rad': targets.tolist(), 'gripper_0_255': jaw} |
|
|
| def step(self, action): |
| if not self.ready or self.done: |
| raise RuntimeError('Reset is required before another task step') |
| if type(action) is not dict or set(action) != {'joint_targets_rad', 'gripper_0_255'}: |
| raise ValueError('Task action requires six joint targets and one native gripper command') |
| mj, np = sim._native_imports() |
| targets = np.asarray(_vector(action['joint_targets_rad'], 6, 'joint targets'), dtype=float) |
| jaw = _number(action['gripper_0_255'], 'gripper command', 0, 255) |
| if np.any(targets < self.model.actuator_ctrlrange[:6, 0]) or np.any(targets > self.model.actuator_ctrlrange[:6, 1]): |
| raise ValueError('Task joint target exceeds native actuator limits') |
| if not self._finite(): |
| self.reason, self.done = 'invalid_state', True |
| return None, -1., True, False, self.info() |
| before_distance = float(np.linalg.norm(self.data.xpos[self.block]-self.config.target_position_m)) |
| messages = [] |
| with sim.native_warnings(messages): |
| for _ in range(round(self.config.control_dt/self.config.timestep)): |
| self.command_targets += np.clip(targets-self.command_targets, |
| -self.config.joint_target_rate_rad_s*self.config.timestep, |
| self.config.joint_target_rate_rad_s*self.config.timestep) |
| self.command_gripper += float(np.clip(jaw-self.command_gripper, |
| -self.config.jaw_target_rate_0_255_s*self.config.timestep, |
| self.config.jaw_target_rate_0_255_s*self.config.timestep)) |
| self._controls() |
| mj.mj_step(self.model, self.data) |
| mj.mj_forward(self.model, self.data) |
| self.elapsed += self.config.timestep |
| if not self._finite(): |
| self.reason, self.done = 'invalid_state', True |
| break |
| if messages or any(x.number for x in self.data.warning): |
| self.reason, self.done = 'solver_warning', True |
| break |
| counts, forces, supported, critical = self._contacts() |
| both = all(counts[side] > 0 for side in ('left', 'right')) |
| lift = float(self.data.xpos[self.block, 2]-self.initial_object_position[2]) |
| velocity = np.empty(6) |
| mj.mj_objectVelocity(self.model, self.data, mj.mjtObj.mjOBJ_BODY, self.block, velocity, 0) |
| distance = float(np.linalg.norm(self.data.xpos[self.block]-self.config.target_position_m)) |
| drivers = self.data.qpos[self.linkage_qpos[[0, 4]]] |
| released = not any(counts.values()) and bool(np.max(drivers) < .1) |
| settled = np.linalg.norm(velocity[3:]) < .035 and np.linalg.norm(velocity[:3]) < .3 |
| clear = self.data.site_xpos[self.tcp, 2] >= self.data.xpos[self.block, 2]+self.config.retreat_clearance_m |
| self._outcome_tick(both=both, supported=supported, lift=lift, distance=distance, |
| released=released, settled=settled, clear=clear, |
| critical=critical, pad_force=max(forces.values())) |
| self.warnings.extend(messages) |
| self.steps += 1 |
| if not self.done and self.steps >= round(self.config.timeout_s/self.config.control_dt): |
| self.reason, self.done = 'timeout', True |
| observation = self.observation() if self._finite() else None |
| info = self.info() |
| reward = before_distance-info['target_distance_m'] if info['target_distance_m'] is not None else -1. |
| return observation, reward, self.done and self.reason != 'timeout', self.done and self.reason == 'timeout', info |
|
|
|
|
| class ReferenceTaskPolicy: |
| """Authored approach/grasp/place baseline, with no learned model or fallback.""" |
| def __init__(self, environment): |
| self.env, self.phase, self.ticks = environment, 'approach', 0 |
| self.grasp_position = None |
| self.held_offset = None |
|
|
| def __call__(self, observation): |
| _, np = sim._native_imports() |
| c = self.env.config |
| obj, tcp, target = (np.asarray(observation[name]) for name in |
| ('object_position_m', 'tcp_position_m', 'target_position_m')) |
| jaw = 0. |
| both = all(observation['grasp_contact_counts'][side] > 0 for side in ('left', 'right')) |
| if self.phase == 'approach': |
| goal = obj+np.asarray([0., 0., .13]) |
| if np.linalg.norm(tcp-goal) < .008: |
| self.grasp_position = obj+np.asarray([0., 0., .005]) |
| self.phase, self.ticks = 'descend', 0 |
| elif self.phase == 'descend': |
| goal = self.grasp_position |
| if np.linalg.norm(tcp-goal) < .004: |
| self.phase, self.ticks = 'close', 0 |
| elif self.phase == 'close': |
| goal, jaw = self.grasp_position, 180. |
| if both and min(observation['pad_normal_forces_n'].values()) > .1 and self.ticks >= 25: |
| self.held_offset = obj-tcp |
| self.phase, self.ticks = 'lift', 0 |
| elif self.phase == 'lift': |
| goal, jaw = self.grasp_position+np.asarray([0., 0., .13]), 180. |
| if self.env.lift_seen and np.linalg.norm(tcp-goal) < .008: |
| self.phase, self.ticks = 'transfer', 0 |
| elif self.phase == 'transfer': |
| goal, jaw = target-self.held_offset+np.asarray([0., 0., .13]), 180. |
| if np.linalg.norm(tcp-goal) < .008: |
| self.phase, self.ticks = 'place', 0 |
| elif self.phase == 'place': |
| goal, jaw = target-self.held_offset, 180. |
| if observation['table_object_contact_count'] > 0 and np.linalg.norm(obj-target) < .02 and np.linalg.norm(tcp-goal) < .006: |
| self.phase, self.ticks = 'release', 0 |
| elif self.phase == 'release': |
| goal = target-self.held_offset |
| if not any(observation['grasp_contact_counts'].values()) and self.ticks >= 35: |
| self.phase, self.ticks = 'retreat', 0 |
| else: |
| goal = target+np.asarray([0., 0., .135]) |
| if np.linalg.norm(tcp-goal) < .008: |
| self.phase = 'verify' |
| self.ticks += 1 |
| return self.env.pose_action(goal.tolist(), c.tool_orientation_wxyz, jaw) |
|
|
|
|
| def run_episode(assets, config=None, *, seed=0, control='placement', arm_tape=None): |
| """Retain every issued decision and native state for one attempted episode.""" |
| config = ReferenceTaskConfig() if config is None else config |
| if not isinstance(config, ReferenceTaskConfig): |
| raise TypeError('config must be ReferenceTaskConfig') |
| if control not in ('placement', 'hold', 'replay_open_jaw'): |
| raise ValueError('Unknown reference task control') |
| started = time.perf_counter() |
| attempt = {'schema': SCHEMA, 'profile': PROFILE, 'scope': SCOPE, 'physical_robot_ready': False, |
| 'seed': seed, 'control': control, 'config': asdict(config), 'timestep_seconds': config.timestep, |
| 'simulated_seconds': 0., 'valid': False, 'success': False, 'reason': 'not_started', |
| 'error': None, 'warnings': [], 'trajectory': [], 'actions': [], 'arm_tape': [], |
| 'metrics': {}, 'timing': {}, 'contract': task_contract()} |
| source_root = Path(__file__).parent |
| attempt['implementation_sha256'] = {name: sha256((source_root / name).read_bytes()).hexdigest() |
| for name in ('reference_task.py', 'reference_sim.py', 'reference_assets.lock.json')} |
| environment, phase, latencies = None, 'validation', [] |
| try: |
| if control == 'replay_open_jaw': |
| if type(arm_tape) is not list or len(arm_tape) != round(config.timeout_s/config.control_dt): |
| raise ValueError('Open-jaw control requires exactly one arm target for every configured control tick') |
| arm_tape = [list(_vector(row, 6, 'replay arm target')) for row in arm_tape] |
| phase = 'setup' |
| attempt['asset_provenance'] = sim.verify_assets(assets) |
| environment = ReferenceTaskEnv(assets, config) |
| observation, info = environment.reset(seed=seed) |
| attempt['compiled_model_sha256'] = sim.model_digest(environment.model) |
| attempt['initial_object_position_m'] = environment.initial_object_position.tolist() |
| attempt['target_position_m'] = list(config.target_position_m) |
| attempt['trajectory'].append(observation) |
| attempt['timing']['setup_seconds'] = time.perf_counter()-started |
| policy = ReferenceTaskPolicy(environment) |
| hold = list(observation['joint_position_rad']) |
| phase = 'rollout' |
| for index in range(round(config.timeout_s/config.control_dt)): |
| began = time.perf_counter() |
| if control == 'placement': |
| action = policy(observation) |
| action_phase = policy.phase |
| else: |
| action = {'joint_targets_rad': hold if control == 'hold' else arm_tape[index], |
| 'gripper_0_255': 0.} |
| action_phase = control |
| record = {'t_seconds': observation['t_seconds'], **action, 'phase': action_phase} |
| attempt['actions'].append(record) |
| attempt['arm_tape'].append(list(action['joint_targets_rad'])) |
| observation['issued_joint_targets_rad'] = list(action['joint_targets_rad']) |
| observation['issued_gripper_0_255'] = action['gripper_0_255'] |
| observation['command'] = record |
| observation, _, terminated, truncated, info = environment.step(action) |
| latencies.append(time.perf_counter()-began) |
| if observation is not None: |
| attempt['trajectory'].append(observation) |
| if terminated or truncated: |
| break |
| attempt.update(valid=info['valid'], success=info['success'], reason=info['reason'], |
| simulated_seconds=info['simulation_time'], warnings=info['warnings'], metrics=info) |
| except Exception as exc: |
| attempt['error'] = {'phase': phase, 'type': type(exc).__name__, 'message': str(exc)[:500]} |
| attempt['reason'] = 'exception' |
| if environment is not None: |
| attempt['warnings'] = list(getattr(environment, 'warnings', [])) |
| attempt['simulated_seconds'] = getattr(environment, 'elapsed', 0.) |
| elapsed = time.perf_counter()-started |
| attempt['timing'].update(episode_wall_seconds=elapsed, |
| control_step_seconds=sim._summary(latencies) if latencies else None, |
| missed_control_boundaries=sum(latency > config.control_dt for latency in latencies), |
| control_boundaries=len(latencies), |
| simulated_seconds_per_wall_second=attempt['simulated_seconds']/elapsed, |
| scope='Wall computation, not real-time sensor-to-actuator latency. Includes setup, policy/IK, native physics, contacts, observations and trace construction; excludes process startup, output I/O and rendering') |
| return attempt |
|
|
|
|
| def run_matched(assets, config=None, *, seed=0): |
| placement = run_episode(assets, config, seed=seed) |
| controls = [run_episode(assets, config, seed=seed, control='replay_open_jaw', arm_tape=placement['arm_tape'])] |
| controls.append(run_episode(assets, config, seed=seed, control='hold')) |
| return [placement, *controls] |
|
|
