"""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} # The upstream keyframe contains only robot qpos; initialize them # explicitly and leave the free block's compiled reset coordinates. 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() # Authored gravity feedforward, applied inside the original bounded # affine position servos. Actual force/velocity are never overwritten. 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: # Release is allowed only once the still-held block actually rests # at its destination; a dropped/tossed block cannot acquire this latch. 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) # mju_subQuat returns angular velocity in the current local frame; # site rotational Jacobians are in world coordinates. 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]