obs_utils.py

观测/state 展平及字段规则

原文件 · 93 行 · 2946 字节

utils/obs_utils.py
SHA-256: 0db7842a922fa542725a2b3890cf0d7f46bfb77cdcd614f5d7294119ca4fc1cf

1import numpy as np2from gym import spaces3import torch456POLICY_OBS_KEYS = ("ego", "tasks", "others")7POLICY_OBS_FEATURE_RANKS = {8    "ego": 1,9    "tasks": 2,10    "others": 2,11    "reconfiguration": 1,12    "gat": 1,13}1415CENTRAL_STATE_KEYS = ("agents", "tasks", "global", "focal")16CENTRAL_STATE_FEATURE_RANKS = {17    "agents": 2,18    "tasks": 2,19    "global": 1,20    "focal": 2,21}222324def _flat_dim_from_shape(shape):25    return int(np.prod(shape))262728def flat_obs_dim(obs_space):29    if obs_space.__class__.__name__ == "Dict":30        return sum(_flat_dim_from_shape(obs_space.spaces[k].shape) for k in POLICY_OBS_KEYS)31    return int(np.prod(obs_space.shape))323334def flat_obs_space(obs_space):35    if obs_space.__class__.__name__ != "Dict":36        return obs_space37    return spaces.Box(low=-np.inf, high=np.inf, shape=(flat_obs_dim(obs_space),), dtype=np.float32)383940def flat_state_space(state_space):41    if state_space.__class__.__name__ != "Dict":42        return state_space43    dim = sum(_flat_dim_from_shape(state_space.spaces[k].shape) for k in CENTRAL_STATE_KEYS)44    return spaces.Box(low=-np.inf, high=np.inf, shape=(dim,), dtype=np.float32)454647def flatten_obs_dict(obs):48    if not isinstance(obs, dict):49        return obs50    keys = [k for k in POLICY_OBS_KEYS if k in obs]51    if "gat" in obs:52        keys.append("gat")5354    use_torch = any(isinstance(obs[k], torch.Tensor) for k in keys)55    device = next((obs[k].device for k in keys if isinstance(obs[k], torch.Tensor)), None)56    parts = []57    for key in keys:58        value = obs[key]59        if use_torch:60            value = value if isinstance(value, torch.Tensor) else torch.as_tensor(value, dtype=torch.float32, device=device)61            value = value.float()62        else:63            value = np.asarray(value, dtype=np.float32)64        feature_rank = POLICY_OBS_FEATURE_RANKS[key]65        prefix_shape = value.shape[:-feature_rank]66        parts.append(value.reshape(*prefix_shape, -1))67    if use_torch:68        return torch.cat(parts, dim=-1)69    return np.concatenate(parts, axis=-1)707172def flatten_state_dict(state):73    if not isinstance(state, dict):74        return state7576    keys = [k for k in CENTRAL_STATE_KEYS if k in state]77    use_torch = any(isinstance(state[k], torch.Tensor) for k in keys)78    device = next((state[k].device for k in keys if isinstance(state[k], torch.Tensor)), None)79    parts = []80    for key in keys:81        value = state[key]82        if use_torch:83            value = value if isinstance(value, torch.Tensor) else torch.as_tensor(84                value, dtype=torch.float32, device=device85            )86            value = value.float()87        else:88            value = np.asarray(value, dtype=np.float32)89        prefix_shape = value.shape[:-CENTRAL_STATE_FEATURE_RANKS[key]]90        parts.append(value.reshape(*prefix_shape, -1))91    if use_torch:92        return torch.cat(parts, dim=-1)93    return np.concatenate(parts, axis=-1)