How to properly create a batch with torch.Tensor

valentin-cnt: <p>Hello!</p>
<p>I have a repo where I have implemented A2C and PPO and it uses the same code to gather a batch of data:</p>
<pre><code class="lang-auto">    states = np.zeros((args.num_steps, args.num_envs) + obversation_shape)
    actions = np.zeros((args.num_steps, args.num_envs) + action_shape)
    rewards = np.zeros((args.num_steps, args.num_envs))
    flags = np.zeros((args.num_steps, args.num_envs))
    state_values = torch.zeros((args.num_steps, args.num_envs)).to(args.device)
    log_probs = torch.zeros((args.num_steps, args.num_envs)).to(args.device)

    num_updates = int(args.total_timesteps // args.num_steps)
    global_step = 0

    state, _ = envs.reset(seed=SEED)
    next_done = np.zeros(args.num_envs)

    for _ in tqdm(range(num_updates)):
        # log_probs, state_values = [], []
        start = time.perf_counter()

        for i in range(args.num_steps):
            global_step += 1
            flags[i] = next_done

            action, log_prob, state_value = agent.get_action(state)

            next_state, reward, terminated, truncated, infos = envs.step(
                action)

            states[i] = state
            actions[i] = action
            rewards[i] = reward
            log_probs[i] = log_prob
            state_values[i] = state_value
            # state_values.append(state_value)
            # log_probs.append(log_prob)

            state = next_state
            next_done = np.logical_or(terminated, truncated)

            if "final_info" not in infos:
                continue

            for info in infos["final_info"]:
                if info is None:
                    continue

                writer.add_scalar("rollout/episodic_return",
                                  info["episode"]["r"], global_step)
                writer.add_scalar("rollout/episodic_length",
                                  info["episode"]["l"], global_step)

        agent.update_policy(
            {
                "states": torch.from_numpy(states).float().to(args.device),
                "actions": torch.from_numpy(actions).float().to(args.device),
                "last_state": torch.from_numpy(next_state).float().to(
                    args.device),
                "last_flag": torch.from_numpy(next_done).float().to(
                    args.device),
                "rewards": torch.from_numpy(rewards).float().to(args.device),
                "flags": torch.from_numpy(flags).float().to(args.device),
                # "state_values": torch.stack(state_values).squeeze(),
                # "log_probs": torch.stack(log_probs).squeeze()
                "state_values": state_values,
                "log_probs": log_probs
            }, global_step)
</code></pre>
<p>The code runs fines with PPO, but with A2C, I have this bug message:</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “/var/home/valentin/Work/rl-gym/main.py”, line 225, in <br>
main()<br>
File “/var/home/valentin/Work/rl-gym/main.py”, line 201, in main<br>
agent.update_policy(<br>
File “/var/home/valentin/Work/rl-gym/rl_gym/agent.py”, line 54, in update_policy<br>
self.algorithm.update_policy(batch, step)<br>
File “/var/home/valentin/Work/rl-gym/rl_gym/algorithm/a2c.py”, line 93, in update_policy<br>
loss.backward()<br>
File “/var/home/valentin/Work/rl-gym/.venv/lib64/python3.10/site-packages/torch/_tensor.py”, line 488, in backward<br>
torch.autograd.backward(<br>
File “/var/home/valentin/Work/rl-gym/.venv/lib64/python3.10/site-packages/torch/autograd/<strong>init</strong>.py”, line 197, in backward<br>
Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to backward through the graph a second time or if</p>
</blockquote>
<p>I identified where does it comes from and, it is actually the way I’m saving the <em>log_probs</em> during the rollout. When I apply the commented instead (stocking into a list), everything works perfectly.<br>
Same happens with the critic values (stored as state_values), when I put them in Tensor it have a segmentation fault error, in a list it works fine.</p>
<p>At the end I fixed the bug by storing  <em>log_probs</em> and <em>state_values</em> into lists, but I don’t understand why I have to do this. It has been difficult to debug, and I still don’t really get why it works that way and not the other way (storing into fixed tensors, which is more efficient).</p>

vmoens: <p>Hey <a class="mention" href="/u/valentin-cnt">@valentin-cnt</a><br>
I think you should collect data using the torch.no_grad and run your loss independently. For PPO and A2C the loss should not backprop gradients through the log-probability collected during inference.<br>
Here’s a couple of implementations in torchrl for you to get a sense of how it’s done</p>
<p>This is a WIP tutorial on PPO but the code works fine</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/tutorials/pull/2156">
  <header class="source">

      <a href="https://github.com/pytorch/tutorials/pull/2156" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/tutorials</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/tutorials/pull/2156" target="_blank" rel="noopener nofollow ugc">[WIP] PPO Tutorial</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>vmoens:ppo_tutorial</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2022-12-14" data-time="20:47:00" data-timezone="UTC">08:47PM - 14 Dec 22 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="1 commits changed 2 files with 393 additions and 0 deletions">
          <a href="https://github.com/pytorch/tutorials/pull/2156/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+393</span>
            <span class="removed">-0</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Introduces a PPO tutorial using TorchRL.</p>
  </div>

  </article>

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

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

<p>Here’s the PPO loss in torchrl. It’s executed on data that do not carry any gradients so to say (more precisely, they’re not part of any differentiable computational graph)</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/rl/blob/main/torchrl/objectives/ppo.py">
  <header class="source">

      <a href="https://github.com/pytorch/rl/blob/main/torchrl/objectives/ppo.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/rl/blob/main/torchrl/objectives/ppo.py" target="_blank" rel="noopener nofollow ugc">pytorch/rl/blob/main/torchrl/objectives/ppo.py</a></h4>


      <pre><code class="lang-py"># Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import math
from typing import Tuple

import torch
from tensordict.tensordict import TensorDict, TensorDictBase
from torch import distributions as d

from torchrl.modules import SafeModule
from torchrl.objectives.utils import distance_loss

from ..modules.tensordict_module import SafeProbabilisticSequential
from .common import LossModule


class PPOLoss(LossModule):
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/rl/blob/main/torchrl/objectives/ppo.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<p>Hope that helps!</p>

valentin-cnt: <p>Hey <a class="mention" href="/u/vmoens">@vmoens</a><br>
Thank you for your answer. I actually do use torch.no_grad for PPO, but when I use it for A2C, it doesn’t learn anything at all. I guess maybe my implentation is wrong, but since it worked out the way I coded it, I don’t know what’s wrong.</p>
<p>Here’s the code:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/a2c.py">
  <header class="source">

      <a href="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/a2c.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/a2c.py" target="_blank" rel="noopener nofollow ugc">valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/a2c.py</a></h4>


      <pre><code class="lang-py">import torch
from torch.nn.functional import mse_loss
from torch.nn.utils import clip_grad_norm_
from rl_gym.algorithm.base import Base
from rl_gym.neuralnet import ActorCriticNet, CNNActorCritic


class A2C(Base):

    def __init__(self, args, obversation_space: tuple, action_space: tuple,
                 writer, is_continuous):

        super(A2C,
              self).__init__(args, obversation_space,
                             action_space if is_continuous else (), writer)

        self.is_continuous = is_continuous

        if args.cnn:
            self._model = CNNActorCritic(
</code></pre>



  This file has been truncated. <a href="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/a2c.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

  <div style="clear: both"></div>
</aside>
<aside class="onebox githubblob" data-onebox-src="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/ppo.py">
  <header class="source">

      <a href="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/ppo.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/ppo.py" target="_blank" rel="noopener nofollow ugc">valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/ppo.py</a></h4>


      <pre><code class="lang-py">import numpy as np
import torch
from torch.nn.functional import mse_loss
from torch.nn.utils.clip_grad import clip_grad_norm_
from rl_gym.algorithm.base import Base
from rl_gym.neuralnet import ActorCriticNet, CNNActorCritic


class PPO(Base):

    def __init__(self, args, obversation_space: tuple, action_space: tuple,
                 writer, is_continuous):

        super(PPO,
              self).__init__(args, obversation_space,
                             action_space if is_continuous else (), writer)

        self.__n_optim = 4
        self.__eps_clip = 0.2
        self.__value_constant = 0.5
</code></pre>



  This file has been truncated. <a href="https://github.com/valentin-cnt/rl-gym/blob/master/rl_gym/algorithm/ppo.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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


vmoens: <p>I think you need to compute the log prob in your loss, not in the data batch.<br>
Check torchrl’s A2C code</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/rl/blob/main/torchrl/objectives/a2c.py">
  <header class="source">

      <a href="https://github.com/pytorch/rl/blob/main/torchrl/objectives/a2c.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/rl/blob/main/torchrl/objectives/a2c.py" target="_blank" rel="noopener nofollow ugc">pytorch/rl/blob/main/torchrl/objectives/a2c.py</a></h4>


      <pre><code class="lang-py"># Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Tuple

import torch
from tensordict.tensordict import TensorDict, TensorDictBase
from torch import distributions as d

from torchrl.modules import SafeModule
from torchrl.modules.tensordict_module import SafeProbabilisticSequential
from torchrl.objectives.common import LossModule
from torchrl.objectives.utils import distance_loss


class A2CLoss(LossModule):
    """TorchRL implementation of the A2C loss.

</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/rl/blob/main/torchrl/objectives/a2c.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<p>Check the way we get the log-prob: it’s done at learning time, not during inference.</p>
<p>Hope that helps</p>

valentin-cnt: <p>Yes, it works. Thank you for answer, it actually makes to keep sense to seperate the data batch and the update data. I understand it better now, thank you very much</p>