add r2g

pytorch · vmoens · Aug 30, 2023 · Mar 31, 2023 · Apr 4, 2023 · Apr 14, 2023
commit e377ae8360028f2233dc86fd852f11c6001782fa
diff --git a/examples/decision_transformer/dt_online.py b/examples/decision_transformer/dt_online.py
@@ -63,6 +63,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         print(td_test)
 
     collector = make_collector(cfg, policy=actor)
+
     replay_buffer = make_replay_buffer(cfg.replay_buffer)
     for data in collector:
         data_view = data.reshape(-1)

diff --git a/examples/decision_transformer/utils.py b/examples/decision_transformer/utils.py
@@ -4,14 +4,15 @@
 
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import CompositeSpec, LazyMemmapStorage, TensorDictReplayBuffer
-from torchrl.data.replay_buffers.samplers import PrioritizedSampler, RandomSampler
+from torchrl.data.replay_buffers.samplers import RandomSampler
 from torchrl.envs import (
     CatFrames,
     EnvCreator,
     NoopResetEnv,
     ObservationNorm,
     ParallelEnv,
     RenameTransform,
+    Reward2GoTransform,
     StepCounter,
     TargetReturn,
     TensorDictPrimer,
@@ -168,12 +169,10 @@ def make_collector(cfg, policy):
 
 
 def make_replay_buffer(rb_cfg):
-    if rb_cfg.prb:
-        sampler = PrioritizedSampler(max_capacity=rb_cfg.capacity, alpha=0.7, beta=0.5)
-    else:
-        sampler = RandomSampler()
+    r2g = Reward2GoTransform(gamma=1.0, out_keys=["return_to_go"])
+    sampler = RandomSampler()
     return TensorDictReplayBuffer(
-        storage=LazyMemmapStorage(rb_cfg.capacity), sampler=sampler
+        storage=LazyMemmapStorage(rb_cfg.capacity), sampler=sampler, transform=r2g
     )