(cfg: Any | None, stage: str, obs_dim: int, actions_num: int)
| 105 | |
| 106 | |
| 107 | def _build_net_config(cfg: Any | None, stage: str, obs_dim: int, actions_num: int) -> dict[str, Any]: |
| 108 | actor_units = list(_get_cfg_value(cfg, "train.network.mlp.units", DEFAULT_ACTOR_UNITS)) |
| 109 | priv_units = list(_get_cfg_value(cfg, "train.network.priv_mlp.units", DEFAULT_PRIV_UNITS)) |
| 110 | priv_dim = int(_get_cfg_value(cfg, "train.ppo.priv_info_dim", DEFAULT_PRIV_DIM)) |
| 111 | priv_info_cfg = bool(_get_cfg_value(cfg, "train.ppo.priv_info", True)) |
| 112 | |
| 113 | if stage == "stage2": |
| 114 | priv_info = True |
| 115 | proprio_adapt = True |
| 116 | else: |
| 117 | priv_info = priv_info_cfg |
| 118 | proprio_adapt = False |
| 119 | |
| 120 | return { |
| 121 | "actor_units": actor_units, |
| 122 | "priv_mlp_units": priv_units, |
| 123 | "actions_num": actions_num, |
| 124 | "input_shape": (obs_dim,), |
| 125 | "priv_info": priv_info, |
| 126 | "proprio_adapt": proprio_adapt, |
| 127 | "priv_info_dim": priv_dim, |
| 128 | "obs_per_step": obs_dim // 3, |
| 129 | } |
| 130 | |
| 131 | |
| 132 | class Stage2ExportWrapper(torch.nn.Module): |
no test coverage detected