| 156 | |
| 157 | |
| 158 | def _save_deploy_meta( |
| 159 | meta_path: Path, |
| 160 | *, |
| 161 | stage: str, |
| 162 | checkpoint: Path, |
| 163 | onnx_path: Path, |
| 164 | cfg_path: Path | None, |
| 165 | obs_dim: int, |
| 166 | action_dim: int, |
| 167 | prop_hist_len: int, |
| 168 | obs_per_step: int, |
| 169 | dynamic_batch: bool, |
| 170 | normalize_baked_in: bool, |
| 171 | policy_rate: float, |
| 172 | chunk_size: int, |
| 173 | n_action_steps: int, |
| 174 | runtime_reference: dict[str, Any], |
| 175 | ) -> None: |
| 176 | meta = { |
| 177 | "export": { |
| 178 | "stage": stage, |
| 179 | "config_yaml": str(cfg_path) if cfg_path is not None else "", |
| 180 | }, |
| 181 | "io_contract": { |
| 182 | "inputs": ( |
| 183 | [ |
| 184 | {"name": "obs", "shape": ["B", obs_dim], "dtype": "float32"}, |
| 185 | {"name": "proprio_hist", "shape": ["B", prop_hist_len, obs_per_step], "dtype": "float32"}, |
| 186 | ] |
| 187 | if stage == "stage2" |
| 188 | else [ |
| 189 | {"name": "obs", "shape": ["B", obs_dim], "dtype": "float32"}, |
| 190 | {"name": "priv_info", "shape": ["B", DEFAULT_PRIV_DIM], "dtype": "float32"}, |
| 191 | ] |
| 192 | ), |
| 193 | "outputs": [{"name": "action", "shape": ["B", action_dim], "dtype": "float32"}], |
| 194 | "action_semantics": "delta", |
| 195 | "action_formula": "cur_targets = prev_targets + (1/24) * action, then clamp to joint limits", |
| 196 | "action_clip": [-1.0, 1.0], |
| 197 | "policy_rate_hz": float(policy_rate), |
| 198 | "chunk_size": int(chunk_size), |
| 199 | "n_action_steps": int(n_action_steps), |
| 200 | "joint_order_right_hand": RIGHT_HAND_JOINT_ORDER, |
| 201 | "dynamic_batch": dynamic_batch, |
| 202 | }, |
| 203 | "normalization": { |
| 204 | "baked_in_onnx": normalize_baked_in, |
| 205 | "source": "checkpoint running_mean_std (+ sa_mean_std for stage2)", |
| 206 | }, |
| 207 | "runtime_reference": runtime_reference, |
| 208 | "deploy_keep_from_config": [ |
| 209 | "train.network.mlp.units", |
| 210 | "train.network.priv_mlp.units", |
| 211 | "train.ppo.priv_info", |
| 212 | "train.ppo.priv_info_dim", |
| 213 | "env_runtime.randomize_scale_list", |
| 214 | "env_runtime.grasp_cache_file", |
| 215 | "env_runtime.usd_path", |