| import json |
| import os |
| import torch |
| from typing import Any |
|
|
| class AlgSolution: |
|
|
| ACTION_SCALE = 0.5 |
| EE_BODY_NAME_CANDIDATES = ("gripper_base", "piper_gripper_base") |
| ARM_JOINT_NAME_CANDIDATES = ( |
| ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"], |
| ["arm_joint1", "arm_joint2", "arm_joint3", "arm_joint4", "arm_joint5", "arm_joint6"], |
| ) |
|
|
| def __init__(self): |
| policy_path = os.path.dirname(os.path.abspath(__file__)) + '/policy.pt' |
| self.device = 'cuda' |
|
|
| self.policy = torch.jit.load(policy_path, map_location=self.device) |
| self.policy.eval() |
|
|
| self.leg_action_dim = 12 |
| self.wheel_action_dim = 4 |
| self.arm_action_dim = 8 |
|
|
| self.leg_joint_indices = list(range(12)) |
| self.wheel_joint_indices = list(range(12, 16)) |
|
|
| self.train_to_env_action_scale = torch.tensor( |
| [ |
| 0.25, 0.5, 0.5, |
| 0.25, 0.5, 0.5, |
| 0.25, 0.5, 0.5, |
| 0.25, 0.5, 0.5, |
| ], |
| device=self.device, |
| dtype=torch.float32, |
| ).view(1, -1) |
|
|
| self.env_to_train_action_scale = torch.tensor( |
| [ |
| 4.0, 2.0, 2.0, |
| 4.0, 2.0, 2.0, |
| 4.0, 2.0, 2.0, |
| 4.0, 2.0, 2.0, |
| ], |
| device=self.device, |
| dtype=torch.float32, |
| ).view(1, -1) |
|
|
| default_schedule = ( |
| "0:1.5:-0.4:0," |
| "2:1.2:0.0:0," |
| "5:1.2:-0.4:0," |
| "7.4:1.2:0.25:0," |
| "8.8:1.2:0.0:0," |
| "10.5:1.1:-0.2:0," |
| "12.5:1.1:0.2:0" |
| ) |
| self.command_schedule = self._parse_command_schedule( |
| os.environ.get("ATEC_SCORE_SCHEDULE", default_schedule) |
| ) |
| self.current_score = 0.0 |
|
|
| self.arm_default_action = torch.zeros( |
| (1, self.arm_action_dim), |
| device=self.device, |
| dtype=torch.float32, |
| ) |
| self.wheel_action = float(os.environ.get("ATEC_WHEEL_ACTION", "0.0")) |
| self.smooth_alpha = float(os.environ.get("ATEC_SMOOTH_ALPHA", "1.0")) |
| self.prev_action_env = None |
| self.step_count = 0 |
| self.best_score = float("-inf") |
| self.best_score_step = 0 |
| self.giveup_min_score = float(os.environ.get("ATEC_GIVEUP_MIN_SCORE", "999.0")) |
| self.giveup_stall_steps = int(os.environ.get("ATEC_GIVEUP_STALL_STEPS", "2500")) |
| self.giveup_target_score = float(os.environ.get("ATEC_GIVEUP_TARGET_SCORE", "999.0")) |
| self.controller_mode = os.environ.get("ATEC_CONTROLLER_MODE", "schedule").strip().lower() |
| self.manual_cmd_file = os.environ.get("ATEC_MANUAL_CMD_FILE", "/tmp/atec_l0_manual_cmd.json") |
| self.manual_cmd_max = self._parse_float_list( |
| os.environ.get("ATEC_MANUAL_CMD_MAX", "1.8,0.8,1.0"), |
| 3, |
| ) or [1.8, 0.8, 1.0] |
| self.manual_cmd_default = self._parse_float_list( |
| os.environ.get("ATEC_MANUAL_CMD_DEFAULT", "0.0,0.0,0.0"), |
| 3, |
| ) or [0.0, 0.0, 0.0] |
| self.manual_zero_action_stop = os.environ.get("ATEC_MANUAL_ZERO_ACTION_STOP", "1") != "0" |
| self.manual_stop_eps = float(os.environ.get("ATEC_MANUAL_STOP_EPS", "1e-4")) |
| self.manual_rough_assist = os.environ.get("ATEC_MANUAL_ROUGH_ASSIST", "1") != "0" |
| self.manual_rough_score_min = float(os.environ.get("ATEC_MANUAL_ROUGH_SCORE_MIN", "1.95")) |
| self.manual_rough_score_max = float(os.environ.get("ATEC_MANUAL_ROUGH_SCORE_MAX", "5.20")) |
| self.manual_rough_vx_min = float(os.environ.get("ATEC_MANUAL_ROUGH_VX_MIN", "0.85")) |
| self.manual_rough_vx_max = float(os.environ.get("ATEC_MANUAL_ROUGH_VX_MAX", "1.08")) |
| self.manual_rough_vy_clip = float(os.environ.get("ATEC_MANUAL_ROUGH_VY_CLIP", "0.12")) |
| self.manual_rough_vy_bias = float(os.environ.get("ATEC_MANUAL_ROUGH_VY_BIAS", "0.0")) |
| self.manual_rough_yaw_clip = float(os.environ.get("ATEC_MANUAL_ROUGH_YAW_CLIP", "0.12")) |
| self.rough_option_moe = os.environ.get("ATEC_ROUGH_OPTION_MOE", "1") != "0" |
| self.rough_stall_steps = int(os.environ.get("ATEC_ROUGH_STALL_STEPS", "120")) |
| self.rough_progress_eps = float(os.environ.get("ATEC_ROUGH_PROGRESS_EPS", "0.035")) |
| self.rough_escalate_steps = int(os.environ.get("ATEC_ROUGH_ESCALATE_STEPS", "260")) |
| self.rough_escape_period = int(os.environ.get("ATEC_ROUGH_ESCAPE_PERIOD", "45")) |
| self.rough_escape_vx = float(os.environ.get("ATEC_ROUGH_ESCAPE_VX", "0.68")) |
| self.rough_escape_vy = float(os.environ.get("ATEC_ROUGH_ESCAPE_VY", "0.08")) |
| self.rough_escape_yaw = float(os.environ.get("ATEC_ROUGH_ESCAPE_YAW", "0.34")) |
| self.rough_escape_yaw_strong = float(os.environ.get("ATEC_ROUGH_ESCAPE_YAW_STRONG", "0.46")) |
| self.rough_score_anchor = 0.0 |
| self.rough_step_anchor = 0 |
| self.last_option_id = 0 |
| self.last_velocity_command = None |
|
|
|
|
| def get_action_spec(self) -> dict[str, dict[str, Any]] | None: |
| return None |
|
|
| def _parse_float_list(self, raw: str, expected_len: int) -> list[float] | None: |
| try: |
| vals = [float(x.strip()) for x in raw.split(",") if x.strip()] |
| except Exception: |
| return None |
| if len(vals) != expected_len: |
| return None |
| return vals |
|
|
| @staticmethod |
| def _clip(value: float, lo: float, hi: float) -> float: |
| return max(lo, min(hi, value)) |
|
|
| def _parse_command_schedule(self, raw: str) -> list[tuple[float, torch.Tensor]]: |
| stages = [] |
| for item in raw.split(","): |
| item = item.strip() |
| if not item: |
| continue |
| parts = item.split(":") |
| if len(parts) != 4: |
| continue |
| score, vx, vy, wz = (float(x) for x in parts) |
| cmd = torch.tensor([vx, vy, wz], device=self.device, dtype=torch.float32).view(1, 3) |
| stages.append((score, cmd)) |
| if not stages: |
| stages.append((0.0, torch.tensor([1.2, -0.2, 0.0], device=self.device).view(1, 3))) |
| stages.sort(key=lambda x: x[0]) |
| return stages |
|
|
| def _apply_rough_option_moe(self, cmd: list[float]) -> list[float]: |
| if ( |
| not self.rough_option_moe |
| or not (self.manual_rough_score_min <= self.current_score < self.manual_rough_score_max) |
| or cmd[0] <= self.manual_stop_eps |
| ): |
| self.last_option_id = 0 |
| self.rough_score_anchor = self.current_score |
| self.rough_step_anchor = self.step_count |
| return cmd |
|
|
| if self.current_score >= self.rough_score_anchor + self.rough_progress_eps: |
| self.rough_score_anchor = self.current_score |
| self.rough_step_anchor = self.step_count |
|
|
| stall_steps = self.step_count - self.rough_step_anchor |
| if stall_steps >= self.rough_stall_steps: |
| phase = 1.0 if ((stall_steps // max(1, self.rough_escape_period)) % 2 == 0) else -1.0 |
| yaw_mag = self.rough_escape_yaw_strong if stall_steps >= self.rough_escalate_steps else self.rough_escape_yaw |
| self.last_option_id = 2 |
| return [ |
| self._clip(self.rough_escape_vx, 0.40, self.manual_cmd_max[0]), |
| self._clip(phase * self.rough_escape_vy, -self.manual_cmd_max[1], self.manual_cmd_max[1]), |
| self._clip(phase * yaw_mag, -self.manual_cmd_max[2], self.manual_cmd_max[2]), |
| ] |
|
|
| self.last_option_id = 1 |
| return [ |
| self._clip(cmd[0], self.manual_rough_vx_min, self.manual_rough_vx_max), |
| self._clip(cmd[1] + self.manual_rough_vy_bias, -self.manual_rough_vy_clip, self.manual_rough_vy_clip), |
| self._clip(cmd[2], -self.manual_rough_yaw_clip, self.manual_rough_yaw_clip), |
| ] |
|
|
| def _resolve_joint_ids(self, candidates: tuple[list[str], ...]) -> list[int]: |
| last_error = None |
| for names in candidates: |
| try: |
| ids, found_names = self.robot.find_joints(names) |
| except ValueError as err: |
| last_error = err |
| continue |
| if len(ids) == len(names): |
| if candidates is self.ARM_JOINT_NAME_CANDIDATES: |
| self.arm_joint_names = list(found_names) |
| return list(ids) |
| raise ValueError( |
| f"Cannot resolve required joints from candidates: {candidates}. Last error: {last_error}" |
| ) |
|
|
| def _resolve_ee_body_name(self) -> str: |
| last_error = None |
| for name in self.EE_BODY_NAME_CANDIDATES: |
| try: |
| body_ids, _ = self.robot.find_bodies(name) |
| except ValueError as err: |
| last_error = err |
| continue |
| if len(body_ids) == 1: |
| return name |
| raise ValueError( |
| f"Cannot resolve EE body from candidates: {self.EE_BODY_NAME_CANDIDATES}. Last error: {last_error}" |
| ) |
|
|
| def _ensure_cartesian_targets(self): |
| self.cartesian_ctrl.reset() |
|
|
| def _compute_arm_overlay_action(self) -> torch.Tensor: |
| self._ensure_cartesian_targets() |
|
|
| arm_jpos_des = self.cartesian_ctrl.compute_base( |
| self.ee_pos_target_b, |
| self.ee_quat_target_b, |
| ) |
|
|
| full_target = self.robot.data.joint_pos.clone() |
| full_target[:, self.arm_ids] = arm_jpos_des |
| full_target[:, self.gripper_ids] = self.gripper_open_pos.repeat(full_target.shape[0], 1) |
|
|
| return (full_target - self.default_joint_pos) / self.ACTION_SCALE |
|
|
| def _get_velocity_commands(self, proprio: torch.Tensor) -> torch.Tensor: |
| """Return fixed velocity commands for policy input.""" |
| num_envs = proprio.shape[0] |
|
|
| if self.controller_mode == "manual_file": |
| cmd = list(self.manual_cmd_default) |
| try: |
| with open(self.manual_cmd_file, "r", encoding="utf-8") as f: |
| data = json.load(f) |
| if isinstance(data, dict): |
| cmd = [ |
| float(data.get("vx", cmd[0])), |
| float(data.get("vy", cmd[1])), |
| float(data.get("yaw", data.get("wz", cmd[2]))), |
| ] |
| elif isinstance(data, list) and len(data) >= 3: |
| cmd = [float(data[0]), float(data[1]), float(data[2])] |
| except Exception: |
| pass |
| clipped = [ |
| self._clip(cmd[0], -self.manual_cmd_max[0], self.manual_cmd_max[0]), |
| self._clip(cmd[1], -self.manual_cmd_max[1], self.manual_cmd_max[1]), |
| self._clip(cmd[2], -self.manual_cmd_max[2], self.manual_cmd_max[2]), |
| ] |
| if self.manual_rough_assist: |
| clipped = self._apply_rough_option_moe(clipped) |
| out = torch.tensor(clipped, device=self.device, dtype=proprio.dtype).view(1, 3) |
| if num_envs > 1: |
| out = out.repeat(num_envs, 1) |
| self.last_velocity_command = out.clone() |
| return out |
|
|
| base_cmd = self.command_schedule[0][1] |
| for score, cmd_i in self.command_schedule: |
| if self.current_score >= score: |
| base_cmd = cmd_i |
| else: |
| break |
| cmd = base_cmd.to(dtype=proprio.dtype, device=self.device) |
| if num_envs > 1: |
| cmd = cmd.repeat(num_envs, 1) |
| if self.manual_rough_assist: |
| cmd_list = [float(cmd[0, 0]), float(cmd[0, 1]), float(cmd[0, 2])] |
| cmd_list = self._apply_rough_option_moe(cmd_list) |
| cmd = torch.tensor(cmd_list, device=self.device, dtype=proprio.dtype).view(1, 3) |
| if num_envs > 1: |
| cmd = cmd.repeat(num_envs, 1) |
| self.last_velocity_command = cmd.clone() |
| return cmd |
|
|
| def _extract_policy_obs(self, obs, action_dim) -> torch.Tensor: |
| proprio = obs["proprio"].to(self.device) |
|
|
| expected_dim = 3 + 3 + 3 + 3 + action_dim + action_dim + action_dim |
|
|
| idx = 0 |
| _base_lin_vel = proprio[:, idx:idx + 3] |
| idx += 3 |
|
|
| base_ang_vel = proprio[:, idx:idx + 3] |
| idx += 3 |
|
|
| _velocity_commands_env = proprio[:, idx:idx + 3] |
| idx += 3 |
|
|
| projected_gravity = proprio[:, idx:idx + 3] |
| idx += 3 |
|
|
| joint_pos_all = proprio[:, idx:idx + action_dim] |
| idx += action_dim |
|
|
| joint_vel_all = proprio[:, idx:idx + action_dim] |
| idx += action_dim |
|
|
| actions_all = proprio[:, idx:idx + action_dim] |
|
|
| joint_pos_leg = joint_pos_all[:, self.leg_joint_indices] |
| joint_vel_leg = joint_vel_all[:, self.leg_joint_indices] |
| actions_env_leg = actions_all[:, self.leg_joint_indices] |
|
|
| actions_train_leg = actions_env_leg * self.env_to_train_action_scale.to(dtype=proprio.dtype) |
| velocity_commands = self._get_velocity_commands(proprio) |
|
|
| policy_obs = torch.cat( |
| [ |
| base_ang_vel * 0.25, |
| projected_gravity, |
| velocity_commands, |
| joint_pos_leg, |
| joint_vel_leg * 0.05, |
| actions_train_leg, |
| ], |
| dim=-1, |
| ) |
|
|
| return policy_obs |
|
|
| def _map_policy_action_to_env_action(self, action_train: torch.Tensor, action_dim: int) -> torch.Tensor: |
| """Map training-time 12D leg action to current env full-body action.""" |
| if action_train.shape[-1] != self.leg_action_dim: |
| raise ValueError( |
| f"Policy output dim mismatch: got {action_train.shape[-1]}, expected {self.leg_action_dim}" |
| ) |
|
|
| num_envs = action_train.shape[0] |
| leg_action_env = action_train * self.train_to_env_action_scale |
|
|
| action_env = torch.zeros( |
| (num_envs, action_dim), |
| device=self.device, |
| dtype=torch.float32, |
| ) |
|
|
| action_env[:, self.leg_joint_indices] = leg_action_env |
| if action_dim >= self.leg_action_dim + self.wheel_action_dim + self.arm_action_dim: |
| action_env[:, self.wheel_joint_indices] = self.wheel_action |
| arm_start = self.leg_action_dim + self.wheel_action_dim |
| else: |
| arm_start = self.leg_action_dim |
| action_env[:, arm_start:arm_start + self.arm_action_dim] = self.arm_default_action.repeat(num_envs, 1) |
|
|
| return action_env |
| |
| def predicts(self, obs, current_score): |
| """Run policy inference and return current-env full-body action.""" |
| self.current_score = float(current_score) |
| self.step_count += 1 |
| if current_score >= self.giveup_target_score: |
| return {'action': [], 'giveup': True} |
| if current_score > self.best_score + 1e-4: |
| self.best_score = current_score |
| self.best_score_step = self.step_count |
| elif ( |
| current_score >= self.giveup_min_score |
| and self.step_count - self.best_score_step > self.giveup_stall_steps |
| ): |
| return {'action': [], 'giveup': True} |
|
|
| proprio = obs["proprio"].to(self.device) |
| action_dim = (int(proprio.shape[-1]) - 12) // 3 |
| policy_obs = self._extract_policy_obs(obs, action_dim) |
| if ( |
| self.controller_mode == "manual_file" |
| and self.manual_zero_action_stop |
| and self.last_velocity_command is not None |
| and torch.max(torch.abs(self.last_velocity_command)) <= self.manual_stop_eps |
| ): |
| action_env = torch.zeros((proprio.shape[0], action_dim), device=self.device, dtype=torch.float32) |
| self.prev_action_env = action_env.clone() |
| return {'action': action_env.cpu().numpy().tolist(), 'giveup': False} |
|
|
| with torch.inference_mode(): |
| action_train = self.policy(policy_obs) |
|
|
| if not isinstance(action_train, torch.Tensor): |
| action_train = torch.as_tensor( |
| action_train, device=self.device, dtype=torch.float32 |
| ) |
|
|
| action_train = action_train.to(device=self.device, dtype=torch.float32) |
|
|
| if action_train.ndim == 1: |
| action_train = action_train.unsqueeze(0) |
|
|
| action_env = self._map_policy_action_to_env_action(action_train, action_dim) |
| if self.smooth_alpha < 1.0: |
| if self.prev_action_env is None or self.prev_action_env.shape != action_env.shape: |
| self.prev_action_env = action_env.clone() |
| else: |
| action_env = self.smooth_alpha * action_env + (1.0 - self.smooth_alpha) * self.prev_action_env |
| self.prev_action_env = action_env.clone() |
| action_env = action_env.cpu().numpy().tolist() |
| return {'action': action_env, 'giveup': False} |
|
|
