Prerequisites#

This tutorial assumes that you have completed thePPO tutorial which givesan overview of the TorchRL components and dependencies, such astensordict.TensorDict andtensordict.nn.TensorDictModules,although it should besufficiently transparent to be understood without a deep understanding ofthese classes.

The__init__ method#

The parent class of all losses isLossModule.As many other components of the library, itsforward() method expectsas input atensordict.TensorDict instance sampled from an experiencereplay buffer, or any similar data structure. Using this format makes itpossible to re-use the module acrossmodalities, or in complex settings where the model needs to read multipleentries for instance. In other words, it allows us to code a loss module thatis oblivious to the data type that is being given to is and that focuses onrunning the elementary steps of the loss function and only those.

To keep the tutorial as didactic as we can, we’ll be displaying each methodof the class independently and we’ll be populating the class at a laterstage.

Let us start with the__init__()method. DDPG aims at solving a control task with a simple strategy:training a policy to output actions that maximize the value predicted bya value network. Hence, our loss module needs to receive two networks in itsconstructor: an actor and a value networks. We expect both of these to beTensorDict-compatible objects, such astensordict.nn.TensorDictModule.Our loss function will need to compute a target value and fit the valuenetwork to this, and generate an action and fit the policy such that itsvalue estimate is maximized.

The crucial step of theLossModule.__init__() method is the call toconvert_to_functional(). This method will extractthe parameters from the module and convert it to a functional module.Strictly speaking, this is not necessary and one may perfectly code allthe losses without it. However, we encourage its usage for the followingreason.

The reason TorchRL does this is that RL algorithms often execute the samemodel with different sets of parameters, called “trainable” and “target”parameters.The “trainable” parameters are those that the optimizer needs to fit. The“target” parameters are usually a copy of the former’s with some time lag(absolute or diluted through a moving average).These target parameters are used to compute the value associated with thenext observation. One the advantages of using a set of target parametersfor the value model that do not match exactly the current configuration isthat they provide a pessimistic bound on the value function being computed.Pay attention to thecreate_target_params keyword argument below: thisargument tells theconvert_to_functional()method to create a set of target parameters in the loss module to be usedfor target value computation. If this is set toFalse (see the actor networkfor instance) thetarget_actor_network_params attribute will still beaccessible but this will just return adetached version of theactor parameters.

Later, we will see how the target parameters should be updated in TorchRL.

fromtensordict.nnimportTensorDictModule,TensorDictSequentialdef_init(self,actor_network:TensorDictModule,value_network:TensorDictModule,)->None:super(type(self),self).__init__()self.convert_to_functional(actor_network,"actor_network",create_target_params=True,)self.convert_to_functional(value_network,"value_network",create_target_params=True,compare_against=list(actor_network.parameters()),)self.actor_in_keys=actor_network.in_keys# Since the value we'll be using is based on the actor and value network,# we put them together in a single actor-critic container.actor_critic=ActorCriticWrapper(actor_network,value_network)self.actor_critic=actor_criticself.loss_function="l2"

The value estimator loss method#

In many RL algorithm, the value network (or Q-value network) is trained basedon an empirical value estimate. This can be bootstrapped (TD(0), lowvariance, high bias), meaningthat the target value is obtained using the next reward and nothing else, ora Monte-Carlo estimate can be obtained (TD(1)) in which case the wholesequence of upcoming rewards will be used (high variance, low bias). Anintermediate estimator (TD(\(\lambda\))) can also be used to compromisebias and variance.TorchRL makes it easy to use one or the other estimator via theValueEstimators Enum class, which containspointers to all the value estimators implemented. Let us define the defaultvalue function here. We will take the simplest version (TD(0)), and show lateron how this can be changed.

fromtorchrl.objectives.utilsimportValueEstimatorsdefault_value_estimator=ValueEstimators.TD0

We also need to give some instructions to DDPG on how to build the valueestimator, depending on the user query. Depending on the estimator provided,we will build the corresponding module to be used at train time:

fromtorchrl.objectives.utilsimportdefault_value_kwargsfromtorchrl.objectives.valueimportTD0Estimator,TD1Estimator,TDLambdaEstimatordefmake_value_estimator(self,value_type:ValueEstimators,**hyperparams):hp=dict(default_value_kwargs(value_type))ifhasattr(self,"gamma"):hp["gamma"]=self.gammahp.update(hyperparams)value_key="state_action_value"ifvalue_type==ValueEstimators.TD1:self._value_estimator=TD1Estimator(value_network=self.actor_critic,**hp)elifvalue_type==ValueEstimators.TD0:self._value_estimator=TD0Estimator(value_network=self.actor_critic,**hp)elifvalue_type==ValueEstimators.GAE:raiseNotImplementedError(f"Value type{value_type} it not implemented for loss{type(self)}.")elifvalue_type==ValueEstimators.TDLambda:self._value_estimator=TDLambdaEstimator(value_network=self.actor_critic,**hp)else:raiseNotImplementedError(f"Unknown value type{value_type}")self._value_estimator.set_keys(value=value_key)

Themake_value_estimator method can but does not need to be called: ifnot, theLossModule will query this method withits default estimator.

The actor loss method#

The central piece of an RL algorithm is the training loss for the actor.In the case of DDPG, this function is quite simple: we just need to computethe value associated with an action computed using the policy and optimizethe actor weights to maximize this value.

When computing this value, we must make sure to take the value parameters outof the graph, otherwise the actor and value loss will be mixed up.For this, thehold_out_params() functioncan be used.

def_loss_actor(self,tensordict,)->torch.Tensor:td_copy=tensordict.select(*self.actor_in_keys)# Get an action from the actor network: since we made it functional, we need to pass the paramswithself.actor_network_params.to_module(self.actor_network):td_copy=self.actor_network(td_copy)# get the value associated with that actionwithself.value_network_params.detach().to_module(self.value_network):td_copy=self.value_network(td_copy)return-td_copy.get("state_action_value")

The value loss method#

We now need to optimize our value network parameters.To do this, we will rely on the value estimator of our class:

fromtorchrl.objectives.utilsimportdistance_lossdef_loss_value(self,tensordict,):td_copy=tensordict.clone()# V(s, a)withself.value_network_params.to_module(self.value_network):self.value_network(td_copy)pred_val=td_copy.get("state_action_value").squeeze(-1)# we manually reconstruct the parameters of the actor-critic, where the first# set of parameters belongs to the actor and the second to the value function.target_params=TensorDict({"module":{"0":self.target_actor_network_params,"1":self.target_value_network_params,}},batch_size=self.target_actor_network_params.batch_size,device=self.target_actor_network_params.device,)withtarget_params.to_module(self.actor_critic):target_value=self.value_estimator.value_estimate(tensordict).squeeze(-1)# Computes the value loss: L2, L1 or smooth L1 depending on `self.loss_function`loss_value=distance_loss(pred_val,target_value,loss_function=self.loss_function)td_error=(pred_val-target_value).pow(2)returnloss_value,td_error,pred_val,target_value

Putting things together in a forward call#

The only missing piece is the forward method, which will glue together thevalue and actor loss, collect the cost values and write them in aTensorDictdelivered to the user.

fromtensordictimportTensorDict,TensorDictBasedef_forward(self,input_tensordict:TensorDictBase)->TensorDict:loss_value,td_error,pred_val,target_value=self.loss_value(input_tensordict,)td_error=td_error.detach()td_error=td_error.unsqueeze(input_tensordict.ndimension())ifinput_tensordict.deviceisnotNone:td_error=td_error.to(input_tensordict.device)input_tensordict.set("td_error",td_error,inplace=True,)loss_actor=self.loss_actor(input_tensordict)returnTensorDict(source={"loss_actor":loss_actor.mean(),"loss_value":loss_value.mean(),"pred_value":pred_val.mean().detach(),"target_value":target_value.mean().detach(),"pred_value_max":pred_val.max().detach(),"target_value_max":target_value.max().detach(),},batch_size=[],)fromtorchrl.objectivesimportLossModuleclassDDPGLoss(LossModule):default_value_estimator=default_value_estimatormake_value_estimator=make_value_estimator__init__=_initforward=_forwardloss_value=_loss_valueloss_actor=_loss_actor

Now that we have our loss, we can use it to train a policy to solve acontrol task.

Transforms#

Now that we have a base environment, we may want to modify its representationto make it more policy-friendly. In TorchRL, transforms are appended to thebase environment in a specializedtorchr.envs.TransformedEnv class.

• It is common in DDPG to rescale the reward using some heuristic value. Wewill multiply the reward by 5 in this example.

• If we are usingdm_control, it is also important to build an interfacebetween the simulator which works with double precision numbers, and ourscript which presumably uses single precision ones. This transformation goesboth ways: when callingenv.step(), our actions will need to berepresented in double precision, and the output will need to be transformedto single precision.TheDoubleToFloat transform does exactly this: thein_keys list refers to the keys that will need to be transformed fromdouble to float, while thein_keys_inv refers to those that need tobe transformed to double before being passed to the environment.

• We concatenate the state keys together using theCatTensorstransform.

• Finally, we also leave the possibility of normalizing the states: we willtake care of computing the normalizing constants later on.

fromtorchrl.envsimport(CatTensors,DoubleToFloat,EnvCreator,InitTracker,ObservationNorm,ParallelEnv,RewardScaling,StepCounter,TransformedEnv,)defmake_transformed_env(env,):"""Apply transforms to the ``env`` (such as reward scaling and state normalization)."""env=TransformedEnv(env)# we append transforms one by one, although we might as well create the# transformed environment using the `env = TransformedEnv(base_env, transforms)`# syntax.env.append_transform(RewardScaling(loc=0.0,scale=reward_scaling))# We concatenate all states into a single "observation_vector"# even if there is a single tensor, it'll be renamed in "observation_vector".# This facilitates the downstream operations as we know the name of the# output tensor.# In some environments (not half-cheetah), there may be more than one# observation vector: in this case this code snippet will concatenate them# all.selected_keys=list(env.observation_spec.keys())out_key="observation_vector"env.append_transform(CatTensors(in_keys=selected_keys,out_key=out_key))# we normalize the states, but for now let's just instantiate a stateless# version of the transformenv.append_transform(ObservationNorm(in_keys=[out_key],standard_normal=True))env.append_transform(DoubleToFloat())env.append_transform(StepCounter(max_frames_per_traj))# We need a marker for the start of trajectories for our Ornstein-Uhlenbeck (OU)# exploration:env.append_transform(InitTracker())returnenv

Parallel execution#

The following helper function allows us to run environments in parallel.Running environments in parallel can significantly speed up the collectionthroughput. When using transformed environment, we need to choose whether wewant to execute the transform individually for each environment, orcentralize the data and transform it in batch. Both approaches are easy tocode:

env=ParallelEnv(lambda:TransformedEnv(GymEnv("HalfCheetah-v4"),transforms),num_workers=4)env=TransformedEnv(ParallelEnv(lambda:GymEnv("HalfCheetah-v4"),num_workers=4),transforms)

To leverage the vectorization capabilities of PyTorch, we adoptthe first method:

defparallel_env_constructor(env_per_collector,transform_state_dict,):ifenv_per_collector==1:defmake_t_env():env=make_transformed_env(make_env())env.transform[2].init_stats(3)env.transform[2].loc.copy_(transform_state_dict["loc"])env.transform[2].scale.copy_(transform_state_dict["scale"])returnenvenv_creator=EnvCreator(make_t_env)returnenv_creatorparallel_env=ParallelEnv(num_workers=env_per_collector,create_env_fn=EnvCreator(lambda:make_env()),create_env_kwargs=None,pin_memory=False,)env=make_transformed_env(parallel_env)# we call `init_stats` for a limited number of steps, just to instantiate# the lazy buffers.env.transform[2].init_stats(3,cat_dim=1,reduce_dim=[0,1])env.transform[2].load_state_dict(transform_state_dict)returnenv# The backend can be ``gym`` or ``dm_control``backend="gym"

Scaling the reward helps us control the signal magnitude for a moreefficient learning.

reward_scaling=5.0

We also define when a trajectory will be truncated. A thousand steps (500 ifframe-skip = 2) is a good number to use for the cheetah task:

max_frames_per_traj=500

Normalization of the observations#

To compute the normalizing statistics, we run an arbitrary number of randomsteps in the environment and compute the mean and standard deviation of thecollected observations. TheObservationNorm.init_stats() method canbe used for this purpose. To get the summary statistics, we create a dummyenvironment and run it for a given number of steps, collect data over a givennumber of steps and compute its summary statistics.

defget_env_stats():"""Gets the stats of an environment."""proof_env=make_transformed_env(make_env())t=proof_env.transform[2]t.init_stats(init_env_steps)transform_state_dict=t.state_dict()proof_env.close()returntransform_state_dict

Normalization stats#

Number of random steps used as for stats computation usingObservationNorm

init_env_steps=5000transform_state_dict=get_env_stats()

Gym has been unmaintained since 2022 and does not support NumPy 2.0 amongst other critical functionality.Please upgrade to Gymnasium, the maintained drop-in replacement of Gym, or contact the authors of your software and request that they upgrade.Users of this version of Gym should be able to simply replace 'import gym' with 'import gymnasium as gym' in the vast majority of cases.See the migration guide at https://gymnasium.farama.org/introduction/migration_guide/ for additional information.

Number of environments in each data collector

env_per_collector=4

We pass the stats computed earlier to normalize the output of ourenvironment:

parallel_env=parallel_env_constructor(env_per_collector=env_per_collector,transform_state_dict=transform_state_dict,)fromtorchrl.dataimportCompositeSpec

Exploration#

The policy is passed into aOrnsteinUhlenbeckProcessModuleexploration module, as suggested in the original paper.Let’s define the number of frames before OU noise reaches its minimum value

annealing_frames=1_000_000actor_model_explore=TensorDictSequential(actor,OrnsteinUhlenbeckProcessModule(spec=actor.spec.clone(),annealing_num_steps=annealing_frames,).to(device),)ifdevice==torch.device("cpu"):actor_model_explore.share_memory()

Replay buffer storage and batch size#

TorchRL replay buffer counts the number of elements along the first dimension.Since we’ll be feeding trajectories to our buffer, we need to adapt the buffersize by dividing it by the length of the sub-trajectories yielded by ourdata collector.Regarding the batch-size, our sampling strategy will consist in samplingtrajectories of lengthtraj_len=200 before selecting sub-trajectoriesor lengthrandom_crop_len=25 on which the loss will be computed.This strategy balances the choice of storing whole trajectories of a certainlength with the need for providing samples with a sufficient heterogeneityto our loss. The following figure shows the dataflow from a collectorthat gets 8 frames in each batch with 2 environments run in parallel,feeds them to a replay buffer that contains 1000 trajectories andsamples sub-trajectories of 2 time steps each.

Let’s start with the number of frames stored in the buffer

defceil_div(x,y):return-x//(-y)buffer_size=1_000_000buffer_size=ceil_div(buffer_size,traj_len)

Prioritized replay buffer is disabled by default

prb=False

We also need to define how many updates we’ll be doing per batch of datacollected. This is known as the update-to-data orUTD ratio:

update_to_data=64

We’ll be feeding the loss with trajectories of length 25:

random_crop_len=25

In the original paper, the authors perform one update with a batch of 64elements for each frame collected. Here, we reproduce the same ratiobut while realizing several updates at each batch collection. Weadapt our batch-size to achieve the same number of update-per-frame ratio:

batch_size=ceil_div(64*frames_per_batch,update_to_data*random_crop_len)replay_buffer=make_replay_buffer(buffer_size=buffer_size,batch_size=batch_size,random_crop_len=random_crop_len,prefetch=3,prb=prb,)

Target network updater#

Target networks are a crucial part of off-policy RL algorithms.Updating the target network parameters is made easy thanks to theHardUpdate andSoftUpdateclasses. They’re built with the loss module as argument, and the update isachieved via a call toupdater.step() at the appropriate location in thetraining loop.

fromtorchrl.objectives.utilsimportSoftUpdatetarget_net_updater=SoftUpdate(loss_module,eps=1-tau)

Optimizer#

Finally, we will use the Adam optimizer for the policy and value network:

fromtorchimportoptimoptimizer_actor=optim.Adam(loss_module.actor_network_params.values(True,True),lr=1e-4,weight_decay=0.0)optimizer_value=optim.Adam(loss_module.value_network_params.values(True,True),lr=1e-3,weight_decay=1e-2)total_collection_steps=total_frames//frames_per_batch