[real] 整理 Python Sim2Real v2
This commit is contained in:
@@ -0,0 +1,243 @@
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def resolve_policy_path(policy_path: Path | str | None, root: Path | None = None) -> Path:
|
||||
"""Prefer ONNX for deployment while keeping the PT file as source/fallback."""
|
||||
root = root or Path(__file__).resolve().parents[1]
|
||||
if policy_path is None:
|
||||
onnx_path = root / "policies" / "model_rough.onnx"
|
||||
pt_path = root / "policies" / "model_rough.pt"
|
||||
return onnx_path if onnx_path.exists() else pt_path
|
||||
|
||||
path = Path(policy_path)
|
||||
if not path.is_absolute():
|
||||
path = root / path
|
||||
if path.suffix.lower() == ".pt":
|
||||
onnx_path = path.with_suffix(".onnx")
|
||||
if onnx_path.exists():
|
||||
return onnx_path
|
||||
return path
|
||||
|
||||
|
||||
class PolicyMLP(nn.Module):
|
||||
def __init__(self, obs_dim: int, action_dim: int):
|
||||
super().__init__()
|
||||
self.register_buffer("obs_mean", torch.zeros(obs_dim))
|
||||
self.register_buffer("obs_std", torch.ones(obs_dim))
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(obs_dim, 512),
|
||||
nn.ELU(),
|
||||
nn.Linear(512, 256),
|
||||
nn.ELU(),
|
||||
nn.Linear(256, 128),
|
||||
nn.ELU(),
|
||||
nn.Linear(128, action_dim),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = (x - self.obs_mean) / torch.clamp(self.obs_std, min=1e-6)
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class OnnxPolicy:
|
||||
def __init__(self, model_path: Path):
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"onnxruntime is required for ONNX policy inference. "
|
||||
"Install it on Orin with `python -m pip install onnxruntime`."
|
||||
) from exc
|
||||
|
||||
opts = ort.SessionOptions()
|
||||
opts.intra_op_num_threads = 1
|
||||
opts.inter_op_num_threads = 1
|
||||
opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
|
||||
self.session = ort.InferenceSession(
|
||||
str(model_path),
|
||||
sess_options=opts,
|
||||
providers=["CPUExecutionProvider"],
|
||||
)
|
||||
self.input_name = self.session.get_inputs()[0].name
|
||||
self.output_name = self.session.get_outputs()[0].name
|
||||
input_shape = self.session.get_inputs()[0].shape
|
||||
output_shape = self.session.get_outputs()[0].shape
|
||||
self.expected_obs_dim = int(input_shape[1]) if len(input_shape) >= 2 and isinstance(input_shape[1], int) else 53
|
||||
self.expected_action_dim = int(output_shape[1]) if len(output_shape) >= 2 and isinstance(output_shape[1], int) else 16
|
||||
self.backend = "onnxruntime"
|
||||
|
||||
self.obs_mean = torch.zeros(self.expected_obs_dim)
|
||||
self.obs_std = torch.ones(self.expected_obs_dim)
|
||||
|
||||
def __call__(self, x: torch.Tensor) -> torch.Tensor:
|
||||
obs = x.detach().cpu().numpy().astype(np.float32, copy=False)
|
||||
action = self.session.run([self.output_name], {self.input_name: obs})[0]
|
||||
return torch.from_numpy(np.asarray(action, dtype=np.float32)).to(x.device)
|
||||
|
||||
|
||||
def load_policy(model_path: Path, device: torch.device):
|
||||
if model_path.suffix.lower() == ".onnx":
|
||||
return OnnxPolicy(model_path)
|
||||
|
||||
checkpoint = torch.load(model_path, map_location=device, weights_only=False)
|
||||
state_dict = checkpoint["actor_state_dict"]
|
||||
|
||||
input_key = "mlp.0.weight" if "mlp.0.weight" in state_dict else "net.0.weight"
|
||||
output_key = "mlp.6.weight" if "mlp.6.weight" in state_dict else "net.6.weight"
|
||||
obs_dim = int(state_dict[input_key].shape[1])
|
||||
action_dim = int(state_dict[output_key].shape[0])
|
||||
|
||||
model = PolicyMLP(obs_dim=obs_dim, action_dim=action_dim)
|
||||
remapped_state_dict: dict[str, torch.Tensor] = {}
|
||||
for key, value in state_dict.items():
|
||||
if key.startswith("mlp."):
|
||||
remapped_state_dict[key.replace("mlp.", "net.")] = value
|
||||
elif key.startswith("net."):
|
||||
remapped_state_dict[key] = value
|
||||
elif key == "obs_normalizer._mean":
|
||||
remapped_state_dict["obs_mean"] = value.squeeze()
|
||||
elif key == "obs_normalizer._var":
|
||||
remapped_state_dict["obs_std"] = torch.sqrt(value.squeeze() + 1e-5)
|
||||
|
||||
model.load_state_dict(remapped_state_dict, strict=False)
|
||||
model.eval()
|
||||
model.to(device)
|
||||
model.expected_obs_dim = obs_dim
|
||||
model.expected_action_dim = action_dim
|
||||
model.backend = "torch"
|
||||
return model
|
||||
|
||||
|
||||
class PolicyRunner:
|
||||
BASE_OBS_DIM = 53
|
||||
DEFAULT_STAND_POSE = np.array(
|
||||
[
|
||||
0.0, 0.9, -1.8,
|
||||
0.0, 0.9, -1.8,
|
||||
0.0, 0.9, -1.8,
|
||||
0.0, 0.9, -1.8,
|
||||
0.0, 0.0, 0.0, 0.0,
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policy_path: Path,
|
||||
device: torch.device | None = None,
|
||||
enable_zero_cmd_suppression: bool = True,
|
||||
hold_zero_command_pose: bool = True,
|
||||
command_release_s: float = 0.35,
|
||||
action_scale: np.ndarray | None = None,
|
||||
zero_cmd_use_yaw_rate: bool = True,
|
||||
clip_obs: float = 100.0,
|
||||
):
|
||||
self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.policy_path = Path(policy_path)
|
||||
self.enable_zero_cmd_suppression = bool(enable_zero_cmd_suppression)
|
||||
self.hold_zero_command_pose = bool(hold_zero_command_pose)
|
||||
self.command_release_s = max(float(command_release_s), 1e-3)
|
||||
self.clip_obs = max(float(clip_obs), 0.0)
|
||||
self.policy_path = resolve_policy_path(self.policy_path)
|
||||
if not self.policy_path.exists():
|
||||
raise FileNotFoundError(f"policy file not found: {self.policy_path}")
|
||||
print(f"[PolicyRunner] device={self.device}, policy={self.policy_path}")
|
||||
self.policy = load_policy(self.policy_path, self.device)
|
||||
if self.policy.expected_obs_dim != self.BASE_OBS_DIM:
|
||||
raise ValueError(
|
||||
f"Unsupported policy obs dim {self.policy.expected_obs_dim}. "
|
||||
f"Current sim2real only supports {self.BASE_OBS_DIM}-D actor observations."
|
||||
)
|
||||
|
||||
self.default_dof_pos = self.DEFAULT_STAND_POSE.copy()
|
||||
self.last_actions = np.zeros(16, dtype=np.float32)
|
||||
|
||||
self.action_scale = np.asarray(
|
||||
action_scale
|
||||
if action_scale is not None
|
||||
else [
|
||||
0.125, 0.25, 0.25,
|
||||
0.125, 0.25, 0.25,
|
||||
0.125, 0.25, 0.25,
|
||||
0.125, 0.25, 0.25,
|
||||
5.0, 5.0, 5.0, 5.0,
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
if self.action_scale.shape != (16,):
|
||||
raise ValueError(f"action_scale must be shape (16,), got {self.action_scale.shape}")
|
||||
|
||||
self.zero_cmd_lin_thresh = 0.05
|
||||
self.zero_cmd_yaw_thresh = 0.05
|
||||
self.zero_yaw_rate_thresh = 0.10
|
||||
self.zero_cmd_use_yaw_rate = bool(zero_cmd_use_yaw_rate)
|
||||
self._command_release_alpha = 0.0
|
||||
print(
|
||||
f"[PolicyRunner] obs_dim={self.policy.expected_obs_dim}, "
|
||||
f"base_obs_dim={self.BASE_OBS_DIM}, history=1, "
|
||||
f"action_dim={self.policy.expected_action_dim}, "
|
||||
f"backend={getattr(self.policy, 'backend', 'unknown')}, "
|
||||
f"clip_obs={self.clip_obs:.1f}, "
|
||||
f"zero_cmd_suppression={self.enable_zero_cmd_suppression}, "
|
||||
f"hold_zero_command_pose={self.hold_zero_command_pose}"
|
||||
)
|
||||
|
||||
def reset(self, prime_obs: np.ndarray | None = None) -> None:
|
||||
self.last_actions = np.zeros(16, dtype=np.float32)
|
||||
self._command_release_alpha = 0.0
|
||||
|
||||
def _is_zero_command(self, command: np.ndarray, base_ang_vel: np.ndarray) -> bool:
|
||||
cmd_is_zero = (
|
||||
np.linalg.norm(command[:2]) < self.zero_cmd_lin_thresh
|
||||
and abs(command[2]) < self.zero_cmd_yaw_thresh
|
||||
)
|
||||
if not self.zero_cmd_use_yaw_rate:
|
||||
return cmd_is_zero
|
||||
return cmd_is_zero and abs(base_ang_vel[2]) < self.zero_yaw_rate_thresh
|
||||
|
||||
def command_activation_metrics(self, command: np.ndarray) -> tuple[float, float]:
|
||||
command = np.asarray(command, dtype=np.float32)
|
||||
planar = float(np.linalg.norm(command[:2]))
|
||||
yaw = float(abs(command[2]))
|
||||
return planar, yaw
|
||||
|
||||
def is_command_active(self, command: np.ndarray) -> bool:
|
||||
planar, yaw = self.command_activation_metrics(command)
|
||||
return planar >= self.zero_cmd_lin_thresh or yaw >= self.zero_cmd_yaw_thresh
|
||||
|
||||
def step(self, obs: np.ndarray, dt: float = 0.02) -> tuple[np.ndarray, np.ndarray]:
|
||||
obs = np.asarray(obs, dtype=np.float32)
|
||||
expected_obs_dim = int(self.policy.expected_obs_dim)
|
||||
if obs.shape[0] != expected_obs_dim:
|
||||
raise ValueError(
|
||||
f"Observation dim mismatch: got {obs.shape[0]}, expected {expected_obs_dim}."
|
||||
)
|
||||
if self.clip_obs > 0.0:
|
||||
obs = np.clip(obs, -self.clip_obs, self.clip_obs).astype(np.float32, copy=False)
|
||||
|
||||
obs_tensor = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
raw_actions = self.policy(obs_tensor).squeeze(0).cpu().numpy()
|
||||
|
||||
raw_actions = np.clip(raw_actions, -10.0, 10.0).astype(np.float32)
|
||||
command = obs[6:9]
|
||||
base_ang_vel = obs[0:3] / 0.25
|
||||
zero_command = self._is_zero_command(command, base_ang_vel)
|
||||
if zero_command:
|
||||
self._command_release_alpha = 0.0
|
||||
if self.hold_zero_command_pose:
|
||||
raw_actions[:] = 0.0
|
||||
elif self.enable_zero_cmd_suppression:
|
||||
raw_actions[12:16] = 0.0
|
||||
raw_actions[:12] *= 0.5
|
||||
else:
|
||||
self._command_release_alpha = min(1.0, self._command_release_alpha + dt / self.command_release_s)
|
||||
raw_actions *= self._command_release_alpha
|
||||
|
||||
self.last_actions = raw_actions.copy()
|
||||
scaled_actions = raw_actions * self.action_scale
|
||||
return scaled_actions, raw_actions
|
||||
Reference in New Issue
Block a user