torchrl.objectives 包¶
TorchRL 提供了一系列可在訓練指令碼中使用的損失。目的是讓損失易於重用/互換,並且具有簡單的簽名。
TorchRL 損失的主要特點是
它們是狀態化物件:它們包含可訓練引數的副本,因此
loss_module.parameters()會給出訓練演算法所需的任何內容。它們遵循
tensordict約定:torch.nn.Module.forward()方法將接收一個 tensordict 作為輸入,其中包含返回損失值所需的所有必要資訊。它們輸出一個
tensordict.TensorDict例項,其中損失值記錄在"loss_<smth>"下,其中smth是描述損失的字串。Tensordict 中的附加鍵可能是在訓練期間記錄的有用指標。
注意
我們返回獨立損失的原因是讓使用者可以為不同的引數集使用不同的最佳化器。損失的求和可以透過以下方式簡單完成
>>> loss_val = sum(loss for key, loss in loss_vals.items() if key.startswith("loss_"))
注意
可以透過查詢 get_stateful_net() 來初始化損失中的引數,它將返回一個可像任何其他模組一樣初始化的網路的狀態化版本。如果修改是就地進行的,它將向下傳遞給使用相同引數集的任何其他模組(在損失內部和外部):例如,從損失中修改 actor_network 引數也將修改收集器中的 actor。如果引數是就地修改的,則可以使用 from_stateful_net() 將損失中的引數重置為新值。
torch.vmap 和隨機性¶
TorchRL 損失模組有很多對 vmap() 的呼叫,以攤銷在迴圈中呼叫多個類似模型的成本,而是將這些操作向量化。 vmap 需要被明確告知如何在生成隨機數時進行處理。為此,需要設定隨機模式,並且該模式必須是以下之一:“error”(預設,處理偽隨機函式時出錯)、“same”(在批次之間複製結果)或“different”(批次的每個元素被單獨處理)。依賴預設設定通常會導致此類錯誤:
>>> RuntimeError: vmap: called random operation while in randomness error mode.
由於 vmap 的呼叫埋藏在損失模組的深處,TorchRL 提供了一個介面,可以透過 loss.vmap_randomness = str_value 從外部設定該 vmap 模式,有關更多資訊,請參閱 vmap_randomness()。
LossModule.vmap_randomness 預設為 “error”(如果沒有檢測到隨機模組),在其他情況下預設為 “different”。預設情況下,只有少數模組被列為隨機,但可以使用 add_random_module() 函式擴充套件該列表。
訓練價值函式¶
TorchRL 提供了一系列 **價值估計器**,例如 TD(0)、TD(1)、TD(\(\lambda\)) 和 GAE。簡而言之,價值估計器是資料(主要是獎勵和完成狀態)和狀態值(即,擬合用於估計狀態值的功能返回的值)的函式。要了解更多關於價值估計器的資訊,請參閱 Sutton 和 Barto 的 RL 入門,特別是關於價值迭代和 TD 學習的章節。它基於資料和代理對映,對跟隨某個狀態或狀態-動作對的折扣回報進行一定程度的(有偏)估計。這些估計器在兩種情況下使用:
訓練價值網路以學習“真實”狀態值(或狀態-動作值)對映,需要一個目標值來擬合它。估計器越好(偏差越小,方差越小),價值網路就會越好,這反過來可以顯著加快策略訓練。通常,價值網路的損失將如下所示:
>>> value = value_network(states) >>> target_value = value_estimator(rewards, done, value_network(next_state)) >>> value_net_loss = (value - target_value).pow(2).mean()
計算用於策略最佳化的“優勢”訊號。優勢是價值估計(來自估計器,即來自“真實”資料)與價值網路的輸出(即其代理)之間的差值。正優勢可以看作是一個訊號,表明策略的實際表現優於預期,從而表明如果將該軌跡作為示例,則有改進的空間。相反,負優勢表示策略的表現不如預期。
事情並不總是像上面的例子那樣簡單,計算價值估計器或優勢的公式可能比這更復雜一些。為了幫助使用者靈活地使用任何一個價值估計器,我們提供了一個簡單的 API 來動態更改它。這裡以 DQN 為例,但所有模組都將遵循類似的結構:
>>> from torchrl.objectives import DQNLoss, ValueEstimators
>>> loss_module = DQNLoss(actor)
>>> kwargs = {"gamma": 0.9, "lmbda": 0.9}
>>> loss_module.make_value_estimator(ValueEstimators.TDLambda, **kwargs)
該 ValueEstimators 類列出了可供選擇的價值估計器。這使得使用者可以輕鬆地依賴自動完成來做出選擇。
|
RL 損失的父類。 |
|
將一個隨機模組新增到將被 |
DQN¶
|
DQN 損失類。 |
|
分散式的 DQN 損失類。 |
DDPG¶
|
DDPG 損失類。 |
SAC¶
|
TorchRL 實現的 SAC 損失。 |
|
離散 SAC 損失模組。 |
REDQ¶
|
REDQ 損失模組。 |
CrossQ¶
|
TorchRL 實現的 CrossQ 損失。 |
IQL¶
|
TorchRL 實現的 IQL 損失。 |
|
TorchRL 實現的離散 IQL 損失。 |
CQL¶
|
TorchRL 實現的連續 CQL 損失。 |
|
TorchRL 實現的離散 CQL 損失。 |
GAIL¶
|
TorchRL 實現的生成對抗模仿學習 (GAIL) 損失。 |
DT¶
|
TorchRL 對線上決策 Transformer 損失的實現。 |
|
TorchRL 對線上決策 Transformer 損失的實現。 |
TD3¶
|
TD3 損失模組。 |
TD3+BC¶
|
TD3+BC 損失模組。 |
PPO¶
|
PPO 損失的父類。 |
|
裁剪 PPO 損失。 |
|
KL 懲罰 PPO 損失。 |
使用 PPO 和多頭動作策略¶
注意
構建多頭策略時要考慮的主要工具有:CompositeDistribution、ProbabilisticTensorDictModule 和 ProbabilisticTensorDictSequential。在處理這些時,建議在指令碼開頭呼叫 tensordict.nn.set_composite_lp_aggregate(False).set(),以指示 CompositeDistribution 不應聚合對數機率,而應將其作為葉節點寫入 tensordict。
在某些情況下,我們有一個單一的優勢值但採取了多個動作。每個動作都有自己的對數機率和形狀。例如,動作空間可以結構化如下:
>>> action_td = TensorDict(
... agents=TensorDict(
... action0=Tensor(batch, n_agents, f0),
... action1=Tensor(batch, n_agents, f1, f2),
... batch_size=torch.Size((batch, n_agents))
... ),
... batch_size=torch.Size((batch,))
... )
其中 f0、f1 和 f2 是任意整數。
請注意,在 TorchRL 中,根 tensordict 的形狀與環境相同(如果環境是批次鎖定的,則與執行的批次環境數量相同)。如果 tensordict 是從緩衝區取樣的,它還將具有重放緩衝區 batch_size 的形狀。儘管 n_agent 維度對於每個動作都通用,但它通常不會出現在根 tensordict 的 batch-size 中(儘管它出現在包含代理特定資料的子 tensordict 中,根據 MARL API)。
這有一個合乎邏輯的原因:代理的數量可能條件了環境的某些但不是全部規範。例如,某些環境對所有代理共享一個完成狀態。在這種情況下,一個更完整的 tensordict 將如下所示:
>>> action_td = TensorDict(
... agents=TensorDict(
... action0=Tensor(batch, n_agents, f0),
... action1=Tensor(batch, n_agents, f1, f2),
... observation=Tensor(batch, n_agents, f3),
... batch_size=torch.Size((batch, n_agents))
... ),
... done=Tensor(batch, 1),
... [...] # etc
... batch_size=torch.Size((batch,))
... )
請注意,done 狀態和 reward 通常會由一個最右側的單例維度括起來。有關此限制的更多資訊,請參閱此 文件部分。
我們動作相對於各自分佈的對數機率可能看起來像這樣:
>>> action_td = TensorDict(
... agents=TensorDict(
... action0_log_prob=Tensor(batch, n_agents),
... action1_log_prob=Tensor(batch, n_agents, f1),
... batch_size=torch.Size((batch, n_agents))
... ),
... batch_size=torch.Size((batch,))
... )
或
>>> action_td = TensorDict(
... agents=TensorDict(
... action0_log_prob=Tensor(batch, n_agents),
... action1_log_prob=Tensor(batch, n_agents),
... batch_size=torch.Size((batch, n_agents))
... ),
... batch_size=torch.Size((batch,))
... )
也就是說,分佈對數機率的維度數量通常從樣本的維度到任何低於它的維度不等,例如,如果分佈是多元的(例如 Dirichlet)或 Independent 例項。相反,tensordict 的維度仍然與 env 的/重放緩衝區 的 batch-size 匹配。
在 PPO 損失呼叫期間,損失模組將按如下方式執行一系列操作:
>>> def ppo(tensordict):
... prev_log_prob = tensordict.select(*log_prob_keys)
... action = tensordict.select(*action_keys)
... new_log_prob = dist.log_prob(action)
... log_weight = new_log_prob - prev_log_prob
... advantage = tensordict.get("advantage") # computed by GAE earlier
... # attempt to map shape
... log_weight.batch_size = advantage.batch_size[:-1]
... log_weight = sum(log_weight.sum(dim="feature").values(True, True)) # get a single tensor of log_weights
... return minimum(log_weight.exp() * advantage, log_weight.exp().clamp(1-eps, 1+eps) * advantage)
要理解 PPO 流水線在多頭策略下的樣子,可以在庫的 示例目錄 中找到一個示例。
A2C¶
|
TorchRL 實現的 A2C 損失。 |
Reinforce¶
|
Reinforce 損失模組。 |
Dreamer¶
|
Dreamer Actor 損失。 |
|
Dreamer Model 損失。 |
|
Dreamer Value 損失。 |
多智慧體目標¶
這些目標特定於多智慧體演算法。
QMixer¶
|
QMixer 損失類。 |
Returns¶
|
價值函式模組的抽象父類。 |
|
優勢函式的時序差分 (TD(0)) 估計。 |
|
優勢函式的無窮時序差分 (TD(1)) 估計。 |
|
優勢函式的 TD(\(\lambda\)) 估計。 |
|
廣義優勢估計函數週圍的類包裝器。 |
|
軌跡的 TD(0) 折扣回報估計。 |
|
軌跡的 TD(0) 優勢估計。 |
|
TD(1) 回報估計。 |
|
向量化 TD(1) 回報估計。 |
|
TD(1) 優勢估計。 |
|
向量化 TD(1) 優勢估計。 |
|
TD(\(\lambda\)) 回報估計。 |
向量化 TD(\(\lambda\)) 回報估計。 |
|
TD(\(\lambda\)) 優勢估計。 |
|
向量化 TD(\(\lambda\)) 優勢估計。 |
|
軌跡的廣義優勢估計。 |
|
軌跡的向量化廣義優勢估計。 |
|
|
計算給定多個軌跡和情節結束的折扣累計獎勵總和。 |
Utils¶
|
Double DQN/DDPG 中用於目標網路更新的硬更新類(與軟更新相對)。 |
|
Double DQN/DDPG 中用於目標網路更新的軟更新類。 |
|
用於自定義構建估計器的價值函式列舉。 |
|
預設價值函式關鍵字引數生成器。 |
|
計算兩個張量之間的距離損失。 |
|
將多個最佳化器組合成一個。 |
|
將網路從計算圖中排除的上下文管理器。 |
|
將引數列表從計算圖中排除的上下文管理器。 |
|
計算下一個狀態值(無梯度),以計算目標值。 |