How to use DataLoader for ReplayBuffer

zer0ne: <p>In RL, the data is not static but keeps growing due to new samples explored by the agent.</p>
<p>I would like to use <code>DataLoader</code> for preparing/loading data from a replay buffer more efficiently. However, it seems that the concept of <code>DataLoader</code> is not well designed for non-stationary data.</p>
<p>So, what would be the best way to extract/load/transform data from a large replay buffer efficiently?</p>
<p>Thanks.</p>

zer0ne: <p>Bringing this thread up.</p>

NotNANtoN: <p>Did you come up with a good solution? I am running into the same problem… The dataloader is usable if you set the number of workers to 0, but that defeats the actual usefulness of a dataloader.</p>
<p>At the moment I am just sampling from large tensors, so I do not have a good solution. I am thinking of using Ray and starting the batch collection process that way.</p>

ptrblck: <p><a href="https://pytorch.org/docs/stable/data.html#iterable-style-datasets"><code>IterableDataset</code></a> was introduced recently, which could be a good candidate for your use case.<br>
I’m not deeply familiar with RL, but based on the descriptions in this topic, it seems the map-style datasets have some limitations in RL setups.</p>

NotNANtoN: <p>Thank you for your suggestion. As a rough overview of what we would need for RL is to sample from a growing dataset. For every batch, newly added samples in the dataset should be considered as well.</p>
<p>The problem I had with the regular map-style approach was that setting the number of workers to &gt; 0 spawns new processes which seem to receive a deep copy of the sampler.  The sampler that they received assumes the dataset size to be constant, which does not work for RL. I thought of somehow fixing this by writing my own sampler, but I do not know if this is going to work. Concretely, I get the error:</p>
<pre><code class="lang-auto">    return iter(torch.randint(high=n, size=(self.num_samples,), dtype=torch.int64).tolist())
RuntimeError: invalid argument 2: max must be greater than min, but got: min = 0, max = 0 at /pytorch/aten/src/TH/generic/THTensorRandom.cpp:43
</code></pre>
<p>So the problem is that self.num_samples stays at 0. Is there a way to let the sampler update its self.num_samples? Or could a custom sampler maybe access the <strong>len</strong> method of the current dataset?</p>
<p>I also looked into the IterableDataset and do not yet quite see how it would be used for RL. It seems to me that the dataset that will be passed to the worker subprocesses will not update their length, but I haven’t tried it yet. I am thinking that maybe I could modify the <strong>iter</strong> method to return the actual length of the current dataset? But then the dataset would only return the indices and obviously I want it to return the states, actions etc and collate them…</p>

NotNANtoN: <p>Edit: Do not use this! I messed up. Because I did not specify the Dataset to be an IterableDataset, the <strong>iter</strong> method was not used. Instead it simply iterated over indices in the dataset instead of sampling random indices.<br>
I still cannot find a way to make it efficiently run with a dataloader. Multiprocessing works, but is 8x slower with 2 workers and 16x slower with 8 workers.</p>
<p><a class="mention" href="/u/zer0ne">@zer0ne</a> <a class="mention" href="/u/ptrblck">@ptrblck</a> I finally managed to get it working. One can use the standard Dataset (not the IterableDataset!) by implementing the methods:</p>
<pre><code class="lang-auto">def __getitem__(self, key):
    pass
    # return single transition: state, action, reward, next_state, done
</code></pre>
<p>and</p>
<pre><code class="lang-auto">    def __iter__(self):
        while True:
            idx = random.randint(0, len(self) - 1)
            yield self[idx]
</code></pre>
<p>Then we need to instantiate this class and create a dataloader based off it. The dataloader can then be used in this way:</p>
<pre><code class="lang-auto">    def sample(self):
        try:
            out = next(self.iter)
        except StopIteration:
            self.iter = iter(self.dataloader)
            out = next(self.iter)
        return out
</code></pre>
<p>This will work for any number of workers!<br>
Of course a custom <strong>collate_fn</strong>  needs to be passed into the Dataloader creation.</p>

Yann: <p>Hello,</p>
<p>I am basically trying to do the same thing, with even more difficulty because I want the collated output to live on CUDA, which is tricky with multiprocessing according to the documentation.</p>
<p>Have you found a way to get this working by any chance?</p>

Kalyan_Garigapati: <p>For those looking to implement simmilar class, torchrl (released in 2022) has ReplayBuffer classes just for above usecase of Dataloader with chaning data.<br>
Docs: <a href="https://pytorch.org/rl/stable/reference/data.html#replay-buffers" rel="noopener nofollow ugc">https://pytorch.org/rl/stable/reference/data.html#replay-buffers</a></p>

vmoens: <p>and we’re also working on concatenating the two together for convenience!</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/rl/pull/2384">
  <header class="source">

      <a href="https://github.com/pytorch/rl/pull/2384" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/rl</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/rl/pull/2384" target="_blank" rel="noopener nofollow ugc">[Feature] Pass replay buffers to SyncDataCollector</a>
      </h4>

    <div class="branches">
      <code>pytorch:gh/vmoens/13/base</code> ← <code>pytorch:gh/vmoens/13/head</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2024-08-09" data-time="15:08:24" data-timezone="UTC">03:08PM - 09 Aug 24 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/vmoens" target="_blank" rel="noopener nofollow ugc">
            <img alt="vmoens" src="https://avatars.githubusercontent.com/u/25529882?v=4" class="onebox-avatar-inline" width="20" height="20">
            vmoens
          </a>
        </div>

        <div class="lines" title="3 commits changed 3 files with 94 additions and 15 deletions">
          <a href="https://github.com/pytorch/rl/pull/2384/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+94</span>
            <span class="removed">-15</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __<span class="show-more-container"><a href="https://github.com/pytorch/rl/pull/2384" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">-&gt;__ #2384</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>You will soon be able to put the RB directly in the collector!</p>