Fix DiscreteDecisionTransformer training

takuseno · Oct 7, 2023 · 468ce52 · 468ce52
1 parent 96e39fb
commit 468ce52
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/d3rlpy/algos/transformer/torch/decision_transformer_impl.py b/d3rlpy/algos/transformer/torch/decision_transformer_impl.py
@@ -130,11 +130,11 @@ def __init__(
     @eval_api
     def predict(self, inpt: TorchTransformerInput) -> torch.Tensor:
         # (1, T, A)
-        probs, _ = self._modules.transformer(
+        _, logits = self._modules.transformer(
             inpt.observations, inpt.actions, inpt.returns_to_go, inpt.timesteps
         )
         # (1, T, A) -> (A,)
-        return probs[0][-1]
+        return logits[0][-1]
 
     def inner_update(
         self, batch: TorchTrajectoryMiniBatch, grad_step: int

diff --git a/reproductions/offline/discrete_decision_transformer.py b/reproductions/offline/discrete_decision_transformer.py
@@ -60,7 +60,7 @@ def main() -> None:
         num_layers=6,
         attn_dropout=0.1,
         embed_dropout=0.1,
-        optim_factory=d3rlpy.models.optimizers.AdamFactory(
+        optim_factory=d3rlpy.models.optimizers.AdamWFactory(
             betas=(0.9, 0.95),
             weight_decay=0.1,
         ),