def disc_rew_preprocess_inputs(
observation_space: gym.Space,
action_space: gym.Space,
state: np.ndarray,
action: np.ndarray,
next_state: np.ndarray,
done: np.ndarray,
device: th.device,
scale: bool = False,
) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
state_th = th.as_tensor(state, device=device)
action_th = th.as_tensor(action, device=device)
next_state_th = th.as_tensor(next_state, device=device)
done_th = th.as_tensor(done, device=device)
del state, action, next_state, done # unused
# preprocess
state_th = preprocessing.preprocess_obs(state_th, observation_space, scale)
action_th = preprocessing.preprocess_obs(action_th, action_space, scale)
next_state_th = preprocessing.preprocess_obs(
next_state_th, observation_space, scale
)
done_th = done_th.to(th.float32)
n_gen = len(state_th)
assert state_th.shape == next_state_th.shape
assert len(action_th) == n_gen
return state_th, action_th, next_state_th, done_th
def _torchify_with_space(self, ndarray: np.ndarray, space: gym.Space,
**kwargs) -> th.Tensor:
tensor = th.as_tensor(ndarray, device=self.discrim.device(), **kwargs)
preprocessed = preprocessing.preprocess_obs(
tensor,
space,
# TODO(sam): can I remove "scale" kwarg in DiscrimNet etc.?
normalize_images=self.discrim.scale,
)
return preprocessed
def extract_features(self, obs: th.Tensor) -> th.Tensor:
"""
Preprocess the observation if needed and extract features.
:param obs:
:return:
"""
assert self.features_extractor is not None, "No features extractor was set"
preprocessed_obs = preprocess_obs(obs, self.observation_space, normalize_images=self.normalize_images)
return self.features_extractor(preprocessed_obs)
def extract_features(self, obs: th.Tensor, device) -> th.Tensor:
"""
Preprocess the observation if needed and extract features.
:param obs:
:return:
"""
assert self.features_extractor is not None, "No features extractor was set"
if self.features_extractor_class == NatureCNN:
preprocessed_obs = preprocess_obs(obs, device, self.observation_space, normalize_images=self.normalize_images)
return self.features_extractor(preprocessed_obs)
else:
preprocessed_obs = preprocess_obs(obs, device, self.observation_space, normalize_images=self.normalize_images) #this is the symbolic graph representation
processed_obs = self.features_extractor(deepcopy(preprocessed_obs)) #this is the embedded graph representation.
if th.isnan(processed_obs.x).any():
print("nan detected in node embeddings")
print(processed_obs.x)
print(preprocessed_obs.x)
print(self.features_extractor.parameters)
return processed_obs,preprocessed_obs #this second argument here is dangerous and will likely break a lot of code. Calling functions just need to ignore the second return value
def extract_features(self, obs: th.Tensor) -> th.Tensor:
assert self.features_extractor is not None, 'No feature extractor was set'
if isinstance(self.features_extractor, DictFeaturesExtractor):
# DictFeaturesExtractors do their own processing of each
# internal observation space
return self.features_extractor(obs)
else:
preprocessed_obs = preprocess_obs(
obs,
self.observation_space,
normalize_images=self.normalize_images)
return self.features_extractor(preprocessed_obs)
def extract_features(self, obs: th.Tensor) -> th.Tensor:
"""
Preprocess the observation if needed and extract features.
:param obs:
:return:
"""
assert self.features_extractor is not None, "No features extractor was set"
preprocessed_obs = preprocess_obs(obs, self.observation_space, normalize_images=self.normalize_images)
obs = self.features_extractor(preprocessed_obs, self.current_robot_id)
# self.set_robot_id(self.all_robot_ids) # reset current robot_id after process
return obs
def extract_features(self, obs: th.Tensor) -> th.Tensor:
"""
Copy and inject sensor adaptors
"""
assert self.features_extractor is not None, "No features extractor was set"
preprocessed_obs = preprocess_obs(
obs,
self.observation_space,
normalize_images=self.normalize_images)
# Inject sensor adaptors
preprocessed_obs = self.pns_sensor_adaptor(preprocessed_obs,
self.current_robot_id)
obs = self.features_extractor(preprocessed_obs)
return obs
def compute_output_shape(observation_space, layers):
"""Compute the size of the output after passing an observation from
`observation_space` through the given `layers`."""
# [None] adds a batch dimension to the random observation
torch_obs = torch.tensor(observation_space.sample()[None])
with torch.no_grad():
sample = preprocess_obs(torch_obs,
observation_space,
normalize_images=True)
for layer in layers:
# forward prop to compute the right size
sample = layer(sample)
# make sure batch axis still matches
assert sample.shape[0] == torch_obs.shape[0]
# return everything else
return sample.shape[1:]
def recursive_preprocess(obs_space_dict: dict, observation_dict: dict,
normalize_images: bool) -> dict:
"""
Recursively iterate through `observation_dict` (a possibly-nested dictionary of actual observations),
preprocess each observation according to the SB3 preprocessing implied by the space corresponding
to the space name in `obs_space_dict` (a dict of actual gym Spaces), and add that to a `processed_observations`
dict. The boolean `normalize_images` governing whether image spaces in general are normalized
:return: A nested dictionary of processed observations
"""
processed_observations = dict()
for space_name, obs_space in obs_space_dict.items():
if isinstance(obs_space, OrderedDict):
processed_observations[space_name] = recursive_preprocess(
obs_space, observation_dict[space_name], normalize_images)
else:
processed_observations[space_name] = preprocess_obs(
observation_dict[space_name],
obs_space,
normalize_images=normalize_images)
return processed_observations
