class CategoricalPgAgent(BaseAgent): def __call__(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value = self.model(*model_inputs) return buffer_to((DistInfo(prob=pi), value), device="cpu") def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value = self.model(*model_inputs) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) agent_info = AgentInfo(dist_info=dist_info, value=value) action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value = self.model(*model_inputs) return value.to("cpu")
class RecurrentCategoricalPgAgentBase(BaseAgent): def __call__(self, observation, prev_action, prev_reward, init_rnn_state, device="cpu"): # Assume init_rnn_state already shaped: [N,B,H] prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to( (observation, prev_action, prev_reward, init_rnn_state), device=self.device) pi, value, next_rnn_state = self.model(*model_inputs) dist_info, value = buffer_to((DistInfo(prob=pi), value), device=device) return dist_info, value, next_rnn_state # Leave rnn_state on device. def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward, device="cpu"): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, rnn_state = self.model(*agent_inputs, self.prev_rnn_state) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) # Model handles None, but Buffer does not, make zeros if needed: prev_rnn_state = self.prev_rnn_state if self.prev_rnn_state is not None else buffer_func( rnn_state, torch.zeros_like) # Transpose the rnn_state from [N,B,H] --> [B,N,H] for storage. # (Special case: model should always leave B dimension in.) prev_rnn_state = buffer_method(prev_rnn_state, "transpose", 0, 1) agent_info = AgentInfoRnn(dist_info=dist_info, value=value, prev_rnn_state=prev_rnn_state) action, agent_info = buffer_to((action, agent_info), device=device) self.advance_rnn_state(rnn_state) # Keep on device. return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward, device="cpu"): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value, _rnn_state = self.model(*agent_inputs, self.prev_rnn_state) return value.to(device)
class CategoricalPgAgent(BaseAgent): """ Agent for policy gradient algorithm using categorical action distribution. Same as ``GausssianPgAgent`` and related classes, except uses ``Categorical`` distribution, and has a different interface to the model (model here outputs discrete probabilities rather than means and log_stds)...maybe could reorganize those interfaces to reduce to one poligy gradient agent class. """ def __call__(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value = self.model(*model_inputs) return buffer_to((DistInfo(prob=pi), value), device="cpu") def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value = self.model(*model_inputs) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) agent_info = AgentInfo(dist_info=dist_info, value=value) action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value = self.model(*model_inputs) return value.to("cpu")
class RecurrentCategoricalPgAgentBase(BaseAgent): def __call__(self, observation, prev_action, prev_reward, init_rnn_state): # Assume init_rnn_state already shaped: [N,B,H] prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to( (observation, prev_action, prev_reward, init_rnn_state), device=self.device) pi, value, next_rnn_state = self.model(*model_inputs) dist_info, value = buffer_to((DistInfo(prob=pi), value), device="cpu") return dist_info, value, next_rnn_state # Leave rnn_state on device. def initialize(self, env_spaces, share_memory=False, global_B=1, obs_stats=None, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, obs_stats=obs_stats, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, rnn_state = self.model(*agent_inputs, self.prev_rnn_state) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) # Model handles None, but Buffer does not, make zeros if needed: prev_rnn_state = self.prev_rnn_state or buffer_func( rnn_state, torch.zeros_like) # Transpose the rnn_state from [N,B,H] --> [B,N,H] for storage. # (Special case: model should always leave B dimension in.) prev_rnn_state = buffer_method(prev_rnn_state, "transpose", 0, 1) agent_info = AgentInfoRnn(dist_info=dist_info, value=value, prev_rnn_state=prev_rnn_state) action, agent_info = buffer_to((action, agent_info), device="cpu") self.advance_rnn_state(rnn_state) # Keep on device. return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def curiosity_step(self, curiosity_type, *args): if curiosity_type == 'icm' or curiosity_type == 'disagreement': observation, next_observation, actions = args actions = self.distribution.to_onehot(actions) curiosity_agent_inputs = buffer_to( (observation, next_observation, actions), device=self.device) agent_curiosity_info = IcmInfo() elif curiosity_type == 'ndigo': observation, prev_actions, actions = args actions = self.distribution.to_onehot(actions) prev_actions = self.distribution.to_onehot(prev_actions) curiosity_agent_inputs = buffer_to( (observation, prev_actions, actions), device=self.device) agent_curiosity_info = NdigoInfo(prev_gru_state=None) elif curiosity_type == 'rnd': next_observation, done = args curiosity_agent_inputs = buffer_to((next_observation, done), device=self.device) agent_curiosity_info = RndInfo() r_int = self.model.curiosity_model.compute_bonus( *curiosity_agent_inputs) r_int, agent_curiosity_info = buffer_to((r_int, agent_curiosity_info), device="cpu") return AgentCuriosityStep(r_int=r_int, agent_curiosity_info=agent_curiosity_info) def curiosity_loss(self, curiosity_type, *args): if curiosity_type == 'icm' or curiosity_type == 'disagreement': observation, next_observation, actions, valid = args actions = self.distribution.to_onehot(actions) actions = actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) curiosity_agent_inputs = buffer_to( (observation, next_observation, actions, valid), device=self.device) inv_loss, forward_loss = self.model.curiosity_model.compute_loss( *curiosity_agent_inputs) losses = (inv_loss.to("cpu"), forward_loss.to("cpu")) elif curiosity_type == 'ndigo': observations, prev_actions, actions, valid = args actions = self.distribution.to_onehot(actions) prev_actions = self.distribution.to_onehot(prev_actions) actions = actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) prev_actions = prev_actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) curiosity_agent_inputs = buffer_to( (observations, prev_actions, actions, valid), device=self.device) forward_loss = self.model.curiosity_model.compute_loss( *curiosity_agent_inputs) losses = (forward_loss.to("cpu")) elif curiosity_type == 'rnd': next_observation, valid = args curiosity_agent_inputs = buffer_to((next_observation, valid), device=self.device) forward_loss = self.model.curiosity_model.compute_loss( *curiosity_agent_inputs) losses = (forward_loss.to("cpu")) return losses @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value, _rnn_state = self.model(*agent_inputs, self.prev_rnn_state) return value.to("cpu")
class RADPgVaeAgent(BaseAgent): """ Agent for policy gradient algorithm using categorical action distribution. Same as ``GausssianPgAgent`` and related classes, except uses ``Categorical`` distribution, and has a different interface to the model (model here outputs discrete probabilities in place of means and log_stds, while both output the value estimate). """ def __init__(self, ModelCls=None, model_kwargs=None, initial_model_state_dict=None, data_augs="", vae_loss_type="l2", vae_beta=1.0, sim_loss_coef=0.1, k_dim=24): super().__init__(ModelCls=ModelCls, model_kwargs=model_kwargs, initial_model_state_dict=initial_model_state_dict) self.data_augs = data_augs self.vae_loss_type = vae_loss_type self.vae_beta = vae_beta self.sim_loss_coef = sim_loss_coef self.k_dim = k_dim def aug_obs(self, observation): # Apply initial augmentations for aug, func in self.augs_funcs.items(): if 'cutout' in aug: observation = func(observation) observation = observation.type( torch.float) # Expect torch.uint8 inputs observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. for aug, func in self.augs_funcs.items(): if 'cutout' in aug: continue else: observation = func(observation) return observation def __call__(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) # This is what needs to modified to apply the augmentation from the data. assert len( observation.shape) == 4, "Observation shape was not length 4" observation_one = self.aug_obs(observation) observation_two = self.aug_obs(observation.detach().clone()) # if hasattr(self.model, "final_act") and self.model.final_act == "tanh": # observation_one = 2*observation_one - 1 # observation_two = 2*observation_two - 1 observation_one, observation_two, prev_action, prev_reward = buffer_to( (observation_one, observation_two, prev_action, prev_reward), device=self.device) pi_one, value_one, latent_one, reconstruction_one = self.model( observation_one, prev_action, prev_reward) pi_two, value_two, latent_two, reconstruction_two = self.model( observation_two, prev_action, prev_reward) bs = 2 * len(observation) if self.vae_loss_type == "l2": recon_loss = (torch.sum( (observation_one - reconstruction_one).pow(2)) + torch.sum( (observation_two - reconstruction_two).pow(2))) / bs # elif self.vae_loss_type == "bce": # recon_loss = (torch.nn.functional.binary_cross_entropy(reconstruction, obs) + # torch.nn.functional.binary_cross_entropy(reconstruction, obs)) / 2 # Calculate the similarity loss mu_one, logsd_one = torch.chunk(latent_one, 2, dim=1) mu_two, logsd_two = torch.chunk(latent_two, 2, dim=1) latent_loss_one = torch.sum(-0.5 * (1 + (2 * logsd_one) - mu_one.pow(2) - (2 * logsd_one).exp())) latent_loss_two = torch.sum(-0.5 * (1 + (2 * logsd_two) - mu_two.pow(2) - (2 * logsd_two).exp())) latent_loss = (latent_loss_one + latent_loss_two) / bs mu_one, mu_two = mu_one[:, :self.k_dim], mu_two[:, :self.k_dim] logvar_one, logvar_two = 2 * logsd_one[:, :self. k_dim], 2 * logsd_two[:, :self. k_dim] # KL divergence between original and augmented. sim_loss = torch.sum(logvar_two - logvar_one + 0.5 * (logvar_one.exp() + (mu_one - mu_two).pow(2)) / logvar_two.exp() - 0.5) / (bs // 2) vae_loss = recon_loss + self.vae_beta * latent_loss + self.sim_loss_coef * sim_loss pi_avg = (pi_one + pi_two) / 2 return buffer_to( (DistInfo(prob=pi_avg), value_one, value_two, vae_loss), device="cpu") def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) self.augs_funcs = OrderedDict() aug_to_func = { 'crop': rad.random_crop, 'crop_horiz': rad.random_crop_horizontile, 'grayscale': rad.random_grayscale, 'cutout': rad.random_cutout, 'cutout_color': rad.random_cutout_color, 'flip': rad.random_flip, 'rotate': rad.random_rotation, 'rand_conv': rad.random_convolution, 'color_jitter': rad.random_color_jitter, 'no_aug': rad.no_aug, } if self.data_augs == "": aug_names = [] else: aug_names = self.data_augs.split('-') for aug_name in aug_names: assert aug_name in aug_to_func, 'invalid data aug string' self.augs_funcs[aug_name] = aug_to_func[aug_name] @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) #observation = observation.type(torch.float) # Expect torch.uint8 inputs #observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. if len(observation.shape) == 3: observation = self.aug_obs(observation.unsqueeze(0)).squeeze(0) else: observation = self.aug_obs(observation) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, latent, reconstruction = self.model(*model_inputs) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) agent_info = AgentInfo(dist_info=dist_info, value=value) action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) #observation = observation.type(torch.float) # Expect torch.uint8 inputs #observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. if len(observation.shape) == 3: observation = self.aug_obs(observation.unsqueeze(0)).squeeze(0) else: observation = self.aug_obs(observation) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value, _latent, _reconstruction = self.model(*model_inputs) return value.to("cpu") @torch.no_grad() def reconstructions(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) observation = observation.type(torch.float) observation = observation.mul_(1. / 255) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, _value, _latent, reconstruction = self.model(*model_inputs) return reconstruction.to("cpu")
class RADPgAgent(BaseAgent): """ Agent for policy gradient algorithm using categorical action distribution. Same as ``GausssianPgAgent`` and related classes, except uses ``Categorical`` distribution, and has a different interface to the model (model here outputs discrete probabilities in place of means and log_stds, while both output the value estimate). """ def __init__(self, ModelCls=None, model_kwargs=None, initial_model_state_dict=None, data_augs="", both_actions=False): super().__init__(ModelCls=ModelCls, model_kwargs=model_kwargs, initial_model_state_dict=initial_model_state_dict) self.data_augs = data_augs self.both_actions = both_actions def aug_obs(self, observation): # Apply initial augmentations for aug, func in self.augs_funcs.items(): if 'cutout' in aug: observation = func(observation) observation = observation.type( torch.float) # Expect torch.uint8 inputs observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. for aug, func in self.augs_funcs.items(): if 'cutout' in aug: continue else: observation = func(observation) return observation def __call__(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) # This is what needs to modified to apply the augmentation from the data. assert len( observation.shape) == 4, "Observation shape was not length 4" augmented = self.aug_obs(observation) observation = observation.type( torch.float) # Expect torch.uint8 inputs observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. augmented, observation, prev_action, prev_reward = buffer_to( (augmented, observation, prev_action, prev_reward), device=self.device) # For visualizing the observations # import matplotlib.pyplot as plt # from torchvision.utils import make_grid # def show_imgs(x,max_display=16): # grid = make_grid(x[:max_display],4).permute(1,2,0).cpu().numpy() # plt.xticks([]) # plt.yticks([]) # plt.imshow(grid) # plt.show() # show_imgs(orig_observation) # show_imgs(observation) aug_pi, aug_value = self.model(augmented, prev_action, prev_reward) if self.both_actions: pi, value = self.model(observation, prev_action, prev_reward) return buffer_to( (DistInfo(prob=aug_pi), DistInfo(prob=pi), aug_value, value), device="cpu") return buffer_to((DistInfo(prob=aug_pi), aug_value), device="cpu") def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) self.augs_funcs = OrderedDict() aug_to_func = { 'crop': rad.random_crop, 'crop_horiz': rad.random_crop_horizontile, 'grayscale': rad.random_grayscale, 'cutout': rad.random_cutout, 'cutout_color': rad.random_cutout_color, 'flip': rad.random_flip, 'rotate': rad.random_rotation, 'rand_conv': rad.random_convolution, 'color_jitter': rad.random_color_jitter, 'no_aug': rad.no_aug, } if self.data_augs == "": aug_names = [] else: aug_names = self.data_augs.split('-') for aug_name in aug_names: assert aug_name in aug_to_func, 'invalid data aug string' self.augs_funcs[aug_name] = aug_to_func[aug_name] @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) observation = observation.type( torch.float) # Expect torch.uint8 inputs observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value = self.model(*model_inputs) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) agent_info = AgentInfo(dist_info=dist_info, value=value) action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) observation = observation.type( torch.float) # Expect torch.uint8 inputs observation = observation.mul_(1. / 255) # From [0-255] to [0-1], in place. model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value = self.model(*model_inputs) return value.to("cpu")
class AtariPgAgent(BaseAgent): """Only doing the feedforward agent for now.""" def __init__( self, ModelCls=AtariPgModel, store_latent=False, state_dict_filename=None, load_conv=False, load_all=False, **kwargs ): super().__init__(ModelCls=ModelCls, **kwargs) self.store_latent = store_latent self.state_dict_filename = state_dict_filename self.load_conv = load_conv self.load_all = load_all assert not (load_all and load_conv) self._act_uniform = False def __call__(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, _ = self.model(*model_inputs) # ignore conv output return buffer_to((DistInfo(prob=pi), value), device="cpu") def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): self.model = self.ModelCls( image_shape=env_spaces.observation.shape, action_size=env_spaces.action.n, **self.model_kwargs ) # Model will have stop_grad inside it. if self.load_conv: logger.log("Agent loading state dict: " + self.state_dict_filename) loaded_state_dict = torch.load(self.state_dict_filename, map_location=torch.device('cpu')) # From UL, saves snapshot: params["algo_state_dict"]["encoder"] loaded_state_dict = loaded_state_dict.get("algo_state_dict", loaded_state_dict) loaded_state_dict = loaded_state_dict.get("encoder", loaded_state_dict) # A bit onerous, but ensures that state dicts match: conv_state_dict = OrderedDict([(k.replace("conv.", "", 1), v) for k, v in loaded_state_dict.items() if k.startswith("conv.")]) self.model.conv.load_state_dict(conv_state_dict) logger.log("Agent loaded CONV state dict.") elif self.load_all: # From RL, saves snapshot: params["agent_state_dict"] loaded_state_dict = torch.load(self.state_dict_filename, map_location=torch.device('cpu')) self.load_state_dict(loaded_state_dict["agent_state_dict"]) logger.log("Agnet loaded FULL state dict.") else: logger.log("Agent NOT loading state dict.") if share_memory: self.model.share_memory() self.shared_model = self.model if self.initial_model_state_dict is not None: raise NotImplementedError self.distribution = Categorical(dim=env_spaces.action.n) self.env_spaces = env_spaces self.share_memory = share_memory @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, conv = self.model(*model_inputs) if self._act_uniform: pi[:] = 1. / pi.shape[-1] # uniform dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) agent_info = AgentInfoConv(dist_info=dist_info, value=value, conv=conv if self.store_latent else None) # Don't write extra data. action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value, _ = self.model(*model_inputs) # Ignore conv out return value.to("cpu") def set_act_uniform(self, act_uniform=True): self._act_uniform = act_uniform
class DmlabPgBaseAgent(BaseAgent): """Only doing the feedforward agent for now.""" def __init__( self, ModelCls=DmlabPgLstmModel, store_latent=False, state_dict_filename=None, load_conv=False, load_all=False, **kwargs ): super().__init__(ModelCls=ModelCls, **kwargs) self.store_latent = store_latent self.state_dict_filename = state_dict_filename self.load_conv = load_conv self.load_all = load_all assert not (load_all and load_conv) self._act_uniform = False def __call__(self, observation, prev_action, prev_reward, init_rnn_state): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to( (observation, prev_action, prev_reward, init_rnn_state), device=self.device ) pi, value, next_rnn_state, _ = self.model(*model_inputs) # Ignore conv out dist_info, value = buffer_to((DistInfo(prob=pi), value), device="cpu") return dist_info, value, next_rnn_state def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): self.model = self.ModelCls( image_shape=env_spaces.observation.shape, output_size=env_spaces.action.n, **self.model_kwargs ) # Model will have stop_grad inside it. if self.load_conv: logger.log("Agent loading state dict: " + self.state_dict_filename) loaded_state_dict = torch.load( self.state_dict_filename, map_location=torch.device("cpu") ) # From UL, saves snapshot: params["algo_state_dict"]["encoder"] loaded_state_dict = loaded_state_dict.get( "algo_state_dict", loaded_state_dict ) loaded_state_dict = loaded_state_dict.get("encoder", loaded_state_dict) # A bit onerous, but ensures that state dicts match: conv_state_dict = OrderedDict( [ (k.replace("conv.", "", 1), v) for k, v in loaded_state_dict.items() if k.startswith("conv.") ] ) self.model.conv.load_state_dict(conv_state_dict) logger.log("Agent loaded CONV state dict.") elif self.load_all: # From RL, saves snapshot: params["agent_state_dict"] loaded_state_dict = torch.load( self.state_dict_filename, map_location=torch.device("cpu") ) self.load_state_dict(loaded_state_dict["agent_state_dict"]) logger.log("Agnet loaded FULL state dict.") else: logger.log("Agent NOT loading state dict.") if share_memory: self.model.share_memory() self.shared_model = self.model if self.initial_model_state_dict is not None: raise NotImplementedError self.distribution = Categorical(dim=env_spaces.action.n) self.env_spaces = env_spaces self.share_memory = share_memory @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to( (observation, prev_action, prev_reward), device=self.device ) pi, value, rnn_state, conv = self.model(*model_inputs, self.prev_rnn_state) if self._act_uniform: pi[:] = 1.0 / pi.shape[-1] # uniform dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) # Model handles None, but Buffer does not, make zeros if needed: prev_rnn_state = self.prev_rnn_state or buffer_func(rnn_state, torch.zeros_like) # Transpose the rnn_state from [N,B,H] --> [B,N,H] for storage. # (Special case: model should always leave B dimension in.) prev_rnn_state = buffer_method(prev_rnn_state, "transpose", 0, 1) agent_info = AgentInfoRnnConv( dist_info=dist_info, value=value, prev_rnn_state=prev_rnn_state, conv=conv if self.store_latent else None, ) # Don't write the extra data. action, agent_info = buffer_to((action, agent_info), device="cpu") self.advance_rnn_state(rnn_state) return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to( (observation, prev_action, prev_reward), device=self.device ) _pi, value, _rnn_state, _conv = self.model(*agent_inputs, self.prev_rnn_state) return value.to("cpu") def set_act_uniform(self, act_uniform=True): self._act_uniform = act_uniform
class VAE(BaseUlAlgorithm): """VAE to predict o_t+k from o_t.""" opt_info_fields = tuple(f for f in OptInfo._fields) # copy def __init__( self, batch_size, learning_rate, replay_filepath, delta_T=0, OptimCls=torch.optim.Adam, optim_kwargs=None, initial_state_dict=None, clip_grad_norm=1000., EncoderCls=EncoderModel, encoder_kwargs=None, latent_size=128, ReplayCls=UlForRlReplayBuffer, activation_loss_coefficient=0.0, learning_rate_anneal=None, # cosine learning_rate_warmup=0, # number of updates VaeHeadCls=VaeHeadModel, hidden_sizes=None, # But maybe use for forward prediction DecoderCls=VaeDecoderModel, decoder_kwargs=None, kl_coeff=1., onehot_action=True, validation_split=0.0, n_validation_batches=0, ): optim_kwargs = dict() if optim_kwargs is None else optim_kwargs encoder_kwargs = dict() if encoder_kwargs is None else encoder_kwargs decoder_kwargs = dict() if decoder_kwargs is None else decoder_kwargs save__init__args(locals()) self.c_e_loss = torch.nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX) assert learning_rate_anneal in [None, "cosine"] self._replay_T = delta_T + 1 def initialize(self, n_updates, cuda_idx=None): self.device = torch.device("cpu") if cuda_idx is None else torch.device( "cuda", index=cuda_idx) examples = self.load_replay() self.encoder = self.EncoderCls( image_shape=examples.observation.shape, latent_size=self.latent_size, # UNUSED **self.encoder_kwargs ) if self.onehot_action: act_dim = self.replay_buffer.samples.action.max() + 1 # discrete only self.distribution = Categorical(act_dim) else: act_shape = self.replay_buffer.samples.action.shape[2:] assert len(act_shape) == 1 act_dim = act_shape[0] self.vae_head = self.VaeHeadCls( latent_size=self.latent_size, action_size=act_dim * self.delta_T, hidden_sizes=self.hidden_sizes, ) self.decoder = self.DecoderCls( latent_size=self.latent_size, **self.decoder_kwargs ) self.encoder.to(self.device) self.vae_head.to(self.device) self.decoder.to(self.device) self.optim_initialize(n_updates) if self.initial_state_dict is not None: self.load_state_dict(self.initial_state_dict) def optimize(self, itr): opt_info = OptInfo(*([] for _ in range(len(OptInfo._fields)))) samples = self.replay_buffer.sample_batch(self.batch_size) if self.lr_scheduler is not None: self.lr_scheduler.step(itr) # Do every itr instead of every epoch self.optimizer.zero_grad() recon_loss, kl_loss, conv_output = self.vae_loss(samples) act_loss = self.activation_loss(conv_output) loss = recon_loss + kl_loss + act_loss loss.backward() if self.clip_grad_norm is None: grad_norm = 0. else: grad_norm = torch.nn.utils.clip_grad_norm_( self.parameters(), self.clip_grad_norm) self.optimizer.step() opt_info.reconLoss.append(recon_loss.item()) opt_info.klLoss.append(kl_loss.item()) opt_info.activationLoss.append(act_loss.item()) opt_info.gradNorm.append(grad_norm.item()) opt_info.convActivation.append( conv_output[0].detach().cpu().view(-1).numpy()) # Keep 1 full one. return opt_info def vae_loss(self, samples): observation = samples.observation[0] # [T,B,C,H,W]->[B,C,H,W] target_observation = samples.observation[self.delta_T] if self.delta_T > 0: action = samples.action[:-1] # [T-1,B,A] don't need the last one if self.onehot_action: action = self.distribution.to_onehot(action) t, b = action.shape[:2] action = action.transpose(1, 0) # [B,T-1,A] action = action.reshape(b, -1) else: action = None observation, target_observation, action = buffer_to( (observation, target_observation, action), device=self.device ) h, conv_out = self.encoder(observation) z, mu, logvar = self.vae_head(h, action) recon_z = self.decoder(z) if target_observation.dtype == torch.uint8: target_observation = target_observation.type(torch.float) target_observation = target_observation.mul_(1 / 255.) b, c, h, w = target_observation.shape recon_losses = F.binary_cross_entropy( input=recon_z.reshape(b * c, h, w), target=target_observation.reshape(b * c, h, w), reduction="none", ) if self.delta_T > 0: valid = valid_from_done(samples.done).type(torch.bool) # [T,B] valid = valid[-1] # [B] valid = valid.to(self.device) else: valid = None # all are valid recon_losses = recon_losses.view(b, c, h, w).sum(dim=(2, 3)) # sum over H,W recon_losses = recon_losses.mean(dim=1) # mean over C (o/w loss is HUGE) recon_loss = valid_mean(recon_losses, valid=valid) # mean over batch kl_losses = 1 + logvar - mu.pow(2) - logvar.exp() kl_losses = kl_losses.sum(dim=-1) # sum over latent dimension kl_loss = -0.5 * valid_mean(kl_losses, valid=valid) # mean over batch kl_loss = self.kl_coeff * kl_loss return recon_loss, kl_loss, conv_out def validation(self, itr): logger.log("Computing validation loss...") val_info = ValInfo(*([] for _ in range(len(ValInfo._fields)))) self.optimizer.zero_grad() for _ in range(self.n_validation_batches): samples = self.replay_buffer.sample_batch(self.batch_size, validation=True) with torch.no_grad(): recon_loss, kl_loss, conv_output = self.vae_loss(samples) val_info.reconLoss.append(recon_loss.item()) val_info.klLoss.append(kl_loss.item()) val_info.convActivation.append( conv_output[0].detach().cpu().view(-1).numpy()) # Keep 1 full one. self.optimizer.zero_grad() logger.log("...validation loss completed.") return val_info def state_dict(self): return dict( encoder=self.encoder.state_dict(), vae_head=self.vae_head.state_dict(), decoder=self.decoder.state_dict(), optimizer=self.optimizer.state_dict(), ) def load_state_dict(self, state_dict): self.encoder.load_state_dict(state_dict["encoder"]) self.vae_head.load_state_dict(state_dict["vae_head"]) self.decoder.load_state_dict(state_dict["decoder"]) self.optimizer.load_state_dict(state_dict["optimizer"]) def parameters(self): yield from self.encoder.parameters() yield from self.vae_head.parameters() yield from self.decoder.parameters() def named_parameters(self): """To allow filtering by name in weight decay.""" yield from self.encoder.named_parameters() yield from self.vae_head.named_parameters() yield from self.decoder.named_parameters() def eval(self): self.encoder.eval() # in case of batch norm self.vae_head.eval() self.decoder.eval() def train(self): self.encoder.train() self.vae_head.train() self.decoder.train()
class RecurrentCategoricalOCAgentBase(OCOptimizerMixin, BaseAgent): def __call__(self, observation, prev_action, prev_reward, sampled_option, prev_option, init_rnn_state, device="cpu"): prev_action = self.distribution.to_onehot(prev_action) prev_option = self.distribution_omega.to_onehot_with_invalid( prev_option) model_inputs = buffer_to((observation, prev_action, prev_reward, prev_option, init_rnn_state, sampled_option), device=self.device) pi, beta, q, pi_omega, next_rnn_state = self.model(*model_inputs[:-1]) return buffer_to( (DistInfo(prob=select_at_indexes(sampled_option, pi)), q, beta, DistInfo(prob=pi_omega)), device=device), next_rnn_state def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) self.distribution_omega = Categorical( dim=self.model_kwargs["option_size"]) @torch.no_grad() def step(self, observation, prev_action, prev_reward, device="cpu"): prev_option_input = self._prev_option if prev_option_input is None: # Hack to extract previous option prev_option_input = torch.full_like(prev_action, -1) prev_action = self.distribution.to_onehot(prev_action) prev_option_input = self.distribution_omega.to_onehot_with_invalid( prev_option_input) model_inputs = buffer_to( (observation, prev_action, prev_reward, prev_option_input), device=self.device) pi, beta, q, pi_omega, rnn_state = self.model(*model_inputs, self.prev_rnn_state) dist_info_omega = DistInfo(prob=pi_omega) new_o, terminations = self.sample_option( beta, dist_info_omega) # Sample terminations and options # Model handles None, but Buffer does not, make zeros if needed: prev_rnn_state = self.prev_rnn_state if self.prev_rnn_state is not None else buffer_func( rnn_state, torch.zeros_like) # Transpose the rnn_state from [N,B,H] --> [B,N,H] for storage. # (Special case: model should always leave B dimension in.) prev_rnn_state = buffer_method(prev_rnn_state, "transpose", 0, 1) dist_info = DistInfo(prob=pi) dist_info_o = DistInfo(prob=select_at_indexes(new_o, pi)) action = self.distribution.sample(dist_info_o) agent_info = AgentInfoOCRnn(dist_info=dist_info, dist_info_o=dist_info_o, q=q, value=(pi_omega * q).sum(-1), termination=terminations, dist_info_omega=dist_info_omega, prev_o=self._prev_option, o=new_o, prev_rnn_state=prev_rnn_state) action, agent_info = buffer_to((action, agent_info), device=device) self.advance_oc_state(new_o) self.advance_rnn_state(rnn_state) return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward, device="cpu"): prev_option_input = self._prev_option if prev_option_input is None: # Hack to extract previous option prev_option_input = torch.full_like(prev_action, -1) prev_action = self.distribution.to_onehot(prev_action) prev_option_input = self.distribution_omega.to_onehot_with_invalid( prev_option_input) agent_inputs = buffer_to( (observation, prev_action, prev_reward, prev_option_input), device=self.device) _pi, beta, q, pi_omega, _rnn_state = self.model( *agent_inputs, self.prev_rnn_state) v = (q * pi_omega).sum( -1 ) # Weight q value by probability of option. Average value if terminal q_prev_o = select_at_indexes(self.prev_option, q) beta_prev_o = select_at_indexes(self.prev_option, beta) value = q_prev_o * (1 - beta_prev_o) + v * beta_prev_o return value.to(device)
class CategoricalOCAgentBase(OCOptimizerMixin, BaseAgent): def __call__(self, observation, prev_action, prev_reward, sampled_option, device="cpu"): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to( (observation, prev_action, prev_reward, sampled_option), device=self.device) pi, beta, q, pi_omega = self.model(*model_inputs[:-1]) return buffer_to( (DistInfo(prob=select_at_indexes(sampled_option, pi)), q, beta, DistInfo(prob=pi_omega)), device=device) def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) self.distribution_omega = Categorical( dim=self.model_kwargs["option_size"]) @torch.no_grad() def step(self, observation, prev_action, prev_reward, device="cpu"): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, beta, q, pi_omega = self.model(*model_inputs) dist_info_omega = DistInfo(prob=pi_omega) new_o, terminations = self.sample_option( beta, dist_info_omega) # Sample terminations and options dist_info = DistInfo(prob=pi) dist_info_o = DistInfo(prob=select_at_indexes(new_o, pi)) action = self.distribution.sample(dist_info_o) agent_info = AgentInfoOC(dist_info=dist_info, dist_info_o=dist_info_o, q=q, value=(pi_omega * q).sum(-1), termination=terminations, dist_info_omega=dist_info_omega, prev_o=self._prev_option, o=new_o) action, agent_info = buffer_to((action, agent_info), device=device) self.advance_oc_state(new_o) return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward, device="cpu"): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, beta, q, pi_omega = self.model(*model_inputs) v = (q * pi_omega).sum( -1 ) # Weight q value by probability of option. Average value if terminal q_prev_o = select_at_indexes(self.prev_option, q) beta_prev_o = select_at_indexes(self.prev_option, beta) value = q_prev_o * (1 - beta_prev_o) + v * beta_prev_o return value.to(device)
class RecurrentCategoricalPgAgentBase(BaseAgent): def __call__(self, observation, prev_action, prev_reward, init_rnn_state): # Assume init_rnn_state already shaped: [N,B,H] prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to( (observation, prev_action, prev_reward, init_rnn_state), device=self.device) pi, value, next_rnn_state = self.model(*model_inputs) dist_info, value = buffer_to((DistInfo(prob=pi), value), device="cpu") return dist_info, value, next_rnn_state # Leave rnn_state on device. def initialize(self, env_spaces, share_memory=False, global_B=1, obs_stats=None, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, obs_stats=obs_stats, env_ranks=env_ranks) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, rnn_state = self.model(*agent_inputs, self.prev_rnn_state) dist_info = DistInfo(prob=pi) if self.dual_model: int_pi, int_value, int_rnn_state = self.model_int( *agent_inputs, self.prev_int_rnn_state) dist_int_info = DistInfo(prob=int_pi) if self._mode == "eval": action = self.distribution.sample(dist_info) else: action = self.distribution.sample(dist_int_info) else: action = self.distribution.sample(dist_info) # Model handles None, but Buffer does not, make zeros if needed: prev_rnn_state = self.prev_rnn_state or buffer_func( rnn_state, torch.zeros_like) # Transpose the rnn_state from [N,B,H] --> [B,N,H] for storage. # (Special case: model should always leave B dimension in.) prev_rnn_state = buffer_method(prev_rnn_state, "transpose", 0, 1) if self.dual_model: prev_int_rnn_state = self.prev_int_rnn_state or buffer_func( int_rnn_state, torch.zeros_like) prev_int_rnn_state = buffer_method(prev_int_rnn_state, "transpose", 0, 1) agent_info = AgentInfoRnnTwin( dist_info=dist_info, value=value, prev_rnn_state=prev_rnn_state, dist_int_info=dist_int_info, int_value=int_value, prev_int_rnn_state=prev_int_rnn_state) else: agent_info = AgentInfoRnn(dist_info=dist_info, value=value, prev_rnn_state=prev_rnn_state) action, agent_info = buffer_to((action, agent_info), device="cpu") self.advance_rnn_state(rnn_state) # Keep on device. if self.dual_model: self.advance_int_rnn_state(int_rnn_state) return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def curiosity_step(self, curiosity_type, *args): curiosity_model = self.model.module.curiosity_model if isinstance( self.model, torch.nn.parallel.DistributedDataParallel ) else self.model.curiosity_model curiosity_step_minibatches = self.model_kwargs[ 'curiosity_step_kwargs']['curiosity_step_minibatches'] T, B = args[0].shape[:2] # either observation or next_observation batch_size = B mb_size = batch_size // curiosity_step_minibatches if curiosity_type in {'icm', 'micm', 'disagreement'}: observation, next_observation, actions = args actions = self.distribution.to_onehot(actions) curiosity_agent_inputs = IcmAgentCuriosityStepInputs( observation=observation, next_observation=next_observation, actions=actions) curiosity_agent_inputs = buffer_to(curiosity_agent_inputs, device=self.device) agent_curiosity_info = IcmInfo() elif curiosity_type == 'ndigo': observation, prev_actions, actions = args actions = self.distribution.to_onehot(actions) prev_actions = self.distribution.to_onehot(prev_actions) curiosity_agent_inputs = NdigoAgentCuriosityStepInputs( observations=observation, prev_actions=prev_actions, actions=actions) curiosity_agent_inputs = buffer_to(curiosity_agent_inputs, device=self.device) agent_curiosity_info = NdigoInfo(prev_gru_state=None) elif curiosity_type == 'rnd': next_observation, done = args curiosity_agent_inputs = RndAgentCuriosityStepInputs( next_observation=next_observation, done=done) curiosity_agent_inputs = buffer_to(curiosity_agent_inputs, device=self.device) agent_curiosity_info = RndInfo() # Need to split the intrinsic reward predictions to several minibatches -- otherwise, we will run out of GPU memory r_ints = [] for idxs in iterate_mb_idxs(batch_size, mb_size, shuffle=False): T_idxs = slice(None) B_idxs = idxs mb_r_int = curiosity_model.compute_bonus( *curiosity_agent_inputs[slice(None), B_idxs]) r_ints.append(mb_r_int) r_int = torch.cat(r_ints, dim=1) r_int, agent_curiosity_info = buffer_to((r_int, agent_curiosity_info), device="cpu") return AgentCuriosityStep(r_int=r_int, agent_curiosity_info=agent_curiosity_info) def curiosity_loss(self, curiosity_type, *args): curiosity_model = self.model.module.curiosity_model if isinstance( self.model, torch.nn.parallel.DistributedDataParallel ) else self.model.curiosity_model if curiosity_type in {'icm', 'micm'}: observation, next_observation, actions, valid = args actions = self.distribution.to_onehot(actions) actions = actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) curiosity_agent_inputs = buffer_to( (observation, next_observation, actions, valid), device=self.device) inv_loss, forward_loss = curiosity_model.compute_loss( *curiosity_agent_inputs) # inv_loss, forward_loss = curiosity_model.compute_loss(*args) losses = (inv_loss.to("cpu"), forward_loss.to("cpu")) elif curiosity_type == 'disagreement': observation, next_observation, actions, valid = args actions = self.distribution.to_onehot(actions) actions = actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) curiosity_agent_inputs = buffer_to( (observation, next_observation, actions, valid), device=self.device) forward_loss = curiosity_model.compute_loss( *curiosity_agent_inputs) losses = (forward_loss.to("cpu")) elif curiosity_type == 'ndigo': observations, prev_actions, actions, valid = args actions = self.distribution.to_onehot(actions) prev_actions = self.distribution.to_onehot(prev_actions) actions = actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) prev_actions = prev_actions.squeeze( ) # ([batch, 1, size]) -> ([batch, size]) curiosity_agent_inputs = buffer_to( (observations, prev_actions, actions, valid), device=self.device) forward_loss = curiosity_model.compute_loss( *curiosity_agent_inputs) losses = (forward_loss.to("cpu")) elif curiosity_type == 'rnd': next_observation, valid = args curiosity_agent_inputs = buffer_to((next_observation, valid), device=self.device) forward_loss = curiosity_model.compute_loss( *curiosity_agent_inputs) losses = (forward_loss.to("cpu")) return losses @torch.no_grad() def value(self, observation, prev_action, prev_reward, ret_int=False): prev_action = self.distribution.to_onehot(prev_action) agent_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) # _pi, value, _rnn_state = self.model(*agent_inputs, self.prev_rnn_state) if ret_int: assert self.dual_model _pi, value, _rnn_state = self.model_int(*agent_inputs, self.prev_int_rnn_state) else: _pi, value, _rnn_state = self.model(*agent_inputs, self.prev_rnn_state) return value.to("cpu")
class CategoricalPgAgent(BaseAgent): """ Agent for policy gradient algorithm using categorical action distribution. Same as ``GaussianPgAgent`` and related classes, except uses ``Categorical`` distribution, and has a different interface to the model (model here outputs discrete probabilities in place of means and log_stds, while both output the value estimate). """ def __call__(self, observation, prev_action, prev_reward, dual=False): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value = (self.model_int if dual else self.model)(*model_inputs) return buffer_to((DistInfo(prob=pi), value), device="cpu") def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None, **kwargs): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks, **kwargs) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) # TODO: need to decide which action to take pi, value = self.model(*model_inputs) int_pi, int_value = self.model_int(*model_inputs) dist_info = DistInfo(prob=pi) if self.dual_model: pi_int, pi_int = self.model_int(*model_inputs) dist_int_info = DistInfo(prob=pi_int) if self._mode == "eval": action = self.distribution.sample(dist_info) else: action = self.distribution.sample(dist_int_info) else: action = self.distribution.sample(dist_info) if self.dual_model: agent_info = AgentInfoTwin(dist_info=dist_info, value=value, dist_int_info=dist_int_info, int_value=int_value) else: agent_info = AgentInfo(dist_info=dist_info, value=value) action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward, ret_int=False): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) if ret_int: assert self.dual_model _pi, value = self.model_int(*model_inputs) else: _pi, value = self.model(*model_inputs) return value.to("cpu")
class CategoricalPgVaeAgent(BaseAgent): """ Agent for policy gradient algorithm using categorical action distribution. Same as ``GausssianPgAgent`` and related classes, except uses ``Categorical`` distribution, and has a different interface to the model (model here outputs discrete probabilities in place of means and log_stds, while both output the value estimate). """ def __init__(self,ModelCls,override, model_kwargs, initial_model_state_dict, **agent_config): super().__init__(ModelCls,model_kwargs,initial_model_state_dict,**agent_config) self.override=override def __call__(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, latent, reconstruction, idx = self.model(*model_inputs) return buffer_to((DistInfo(prob=pi), value, latent, reconstruction), device="cpu"), idx def initialize(self, env_spaces, share_memory=False, global_B=1, env_ranks=None): super().initialize(env_spaces, share_memory, global_B=global_B, env_ranks=env_ranks) if self.override['override_policy_value']: policy_layers=self.override["policy_layers"] value_layers=self.override["value_layers"] self.model.override_policy_value(policy_layers=policy_layers, value_layers=value_layers) self.distribution = Categorical(dim=env_spaces.action.n) @torch.no_grad() def step(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) pi, value, latent, reconstruction, idx = self.model(*model_inputs) dist_info = DistInfo(prob=pi) action = self.distribution.sample(dist_info) agent_info = AgentInfoVae(dist_info=dist_info, value=value, latent=latent, reconstruction=reconstruction) action, agent_info = buffer_to((action, agent_info), device="cpu") return AgentStep(action=action, agent_info=agent_info) @torch.no_grad() def value(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, value, _latent, _reconstruction, _idx = self.model(*model_inputs) return value.to("cpu") @torch.no_grad() def reconstruction(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, _value, _latent, reconstruction, _idx = self.model(*model_inputs) return reconstruction.to("cpu") @torch.no_grad() def latent(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, _value, latent, _reconstruction, _idx = self.model(*model_inputs) return latent.to("cpu") @torch.no_grad() def idx(self, observation, prev_action, prev_reward): prev_action = self.distribution.to_onehot(prev_action) model_inputs = buffer_to((observation, prev_action, prev_reward), device=self.device) _pi, _value, latent, _reconstruction, idx = self.model(*model_inputs) return idx