# Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # from dataclasses import dataclass, MISSING from typing import Dict, Iterable, Tuple, Type from benchmarl.algorithms.common import Algorithm, AlgorithmConfig from benchmarl.models.common import ModelConfig from tensordict import TensorDictBase from tensordict.nn import TensorDictModule, TensorDictSequential from torchrl.data import CompositeSpec, UnboundedContinuousTensorSpec from torchrl.modules import EGreedyModule, QValueModule from torchrl.objectives import DQNLoss, LossModule, ValueEstimators class CustomAlgorithm(Algorithm): def __init__( self, delay_value: bool, loss_function: str, my_custom_arg: int, **kwargs ): # In the init function you can define the init parameters you need, just make sure # to pass the kwargs to the super() class super().__init__(**kwargs) self.delay_value = delay_value self.loss_function = loss_function self.my_custom_arg = my_custom_arg # In all the class you have access to a lot of extra things like self.my_custom_method() # Custom methods _ = self.experiment # Experiment class _ = self.experiment_config # Experiment config _ = self.model_config # Policy config _ = self.critic_model_config # Eventual critic config _ = self.group_map # The group to agent names map # Specs _ = self.observation_spec _ = self.action_spec _ = self.state_spec _ = self.action_mask_spec ############################# # Overridden abstract methods ############################# def _get_loss( self, group: str, policy_for_loss: TensorDictModule, continuous: bool ) -> Tuple[LossModule, bool]: if continuous: raise NotImplementedError( "Custom Iql is not compatible with continuous actions." ) else: # Loss loss_module = DQNLoss( policy_for_loss, delay_value=self.delay_value, loss_function=self.loss_function, action_space=self.action_spec[group, "action"], ) # Always tell the loss where to finc the data # You can make sure the data is in the right place in self.process_batch # This loss for example expects all data to have the multagent dimension so we take care of that in # self.process_batch loss_module.set_keys( reward=(group, "reward"), action=(group, "action"), done=(group, "done"), terminated=(group, "terminated"), action_value=(group, "action_value"), value=(group, "chosen_action_value"), priority=(group, "td_error"), ) # Choose your value estimator, see what is available in the ValueEstimators enum loss_module.make_value_estimator( ValueEstimators.TD0, gamma=self.experiment_config.gamma ) # This loss has target delayed parameters so the second value is True return loss_module, True def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]: # For each loss name, associate it the parameters you want # You can optionally modify (aggregate) loss names in self.process_loss_vals() return {"loss": loss.parameters()} def _get_policy_for_loss( self, group: str, model_config: ModelConfig, continuous: bool ) -> TensorDictModule: if continuous: raise ValueError("This should never happen") # The number of agents in the group n_agents = len(self.group_map[group]) # The shape of the discrete action logits_shape = [ *self.action_spec[group, "action"].shape, self.action_spec[group, "action"].space.n, ] # This is the spec of the policy input for this group actor_input_spec = CompositeSpec( {group: self.observation_spec[group].clone().to(self.device)} ) # This is the spec of the policy output for this group actor_output_spec = CompositeSpec( { group: CompositeSpec( {"action_value": UnboundedContinuousTensorSpec(shape=logits_shape)}, shape=(n_agents,), ) } ) # This is our neural policy actor_module = model_config.get_model( input_spec=actor_input_spec, output_spec=actor_output_spec, agent_group=group, input_has_agent_dim=True, # Always true for a policy n_agents=n_agents, centralised=False, # Always false for a policy share_params=self.experiment_config.share_policy_params, device=self.device, action_spec=self.action_spec, ) if self.action_mask_spec is not None: action_mask_key = (group, "action_mask") else: action_mask_key = None value_module = QValueModule( action_value_key=(group, "action_value"), action_mask_key=action_mask_key, out_keys=[ (group, "action"), (group, "action_value"), (group, "chosen_action_value"), ], spec=self.action_spec[group, "action"], action_space=None, # We already passed the spec ) # Here we chain the actor and the value module to get our policy return TensorDictSequential(actor_module, value_module) def _get_policy_for_collection( self, policy_for_loss: TensorDictModule, group: str, continuous: bool ) -> TensorDictModule: if self.action_mask_spec is not None: action_mask_key = (group, "action_mask") else: action_mask_key = None # Add exploration for collection greedy = EGreedyModule( annealing_num_steps=self.experiment_config.get_exploration_anneal_frames( self.on_policy ), action_key=(group, "action"), spec=self.action_spec[(group, "action")], action_mask_key=action_mask_key, eps_init=self.experiment_config.exploration_eps_init, eps_end=self.experiment_config.exploration_eps_end, ) return TensorDictSequential(*policy_for_loss, greedy) def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase: # Here we make sure that all entries have the desired shape, # thus, in case there are shared dones, terminated, or rewards, we expande them keys = list(batch.keys(True, True)) group_shape = batch.get(group).shape nested_done_key = ("next", group, "done") nested_terminated_key = ("next", group, "terminated") nested_reward_key = ("next", group, "reward") if nested_done_key not in keys: batch.set( nested_done_key, batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)), ) if nested_terminated_key not in keys: batch.set( nested_terminated_key, batch.get(("next", "terminated")) .unsqueeze(-1) .expand((*group_shape, 1)), ) if nested_reward_key not in keys: batch.set( nested_reward_key, batch.get(("next", "reward")).unsqueeze(-1).expand((*group_shape, 1)), ) return batch def process_loss_vals( self, group: str, loss_vals: TensorDictBase ) -> TensorDictBase: # Here you can modify the loss_vals tensordict containing entries loss_name->loss_value # For example you can sum two entries in a new entry to optimize them together. return loss_vals ##################### # Custom new methods ##################### def my_custom_method(self): pass @dataclass class CustomAlgorithmConfig(AlgorithmConfig): # This is a class representing the configuration of your algorithm # It will be used to validate loaded configs, so that everytime you load this algorithm # we know exactly which and what parameters to expect with their types # This is a list of args passed to your algorithm delay_value: bool = MISSING loss_function: str = MISSING my_custom_arg: int = MISSING @staticmethod def associated_class() -> Type[Algorithm]: # The associated algorithm class return CustomAlgorithm @staticmethod def supports_continuous_actions() -> bool: # Is it compatible with continuous actions? return False @staticmethod def supports_discrete_actions() -> bool: # Is it compatible with discrete actions? return True @staticmethod def on_policy() -> bool: # Should it be trained on or off policy? return False