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