def build_mdp_features(state, joint_actions):
"""Features used for environment models"""
coordinates = PursuitState.features(state)
actions_one_hot = actions_one_hot_encoding(joint_actions)
# St + At (one hot encoded)
mdp_features = np.concatenate((coordinates, actions_one_hot))
return mdp_features
def extract_features(self, state: PursuitState):
if self._feature_extraction_mode == "default":
return state.features()
elif self._feature_extraction_mode == "relative agent":
return state.features_relative_agent(agent_id=0)
elif self._feature_extraction_mode == "relative prey":
return state.features_relative_prey()
else:
raise ValueError(
f"Invalid feature extraction mode {self._feature_extraction_mode}"
)