Training converges on cpu but never on gpu

username: ozhicha
cooked: <p>Hi,<br>
for the script below my training always converges on CPU and solves the environment. However if I switch DEVICE to cuda it actually never converges.</p>
<p>my python env: torch 2.5.1, torchrl 0.6.0, tensordict 0.6.2</p>
<p>What I tried so far:</p>
<ul>
<li>I also tried running the code on the nightly version of torchrl and the behaviour stayed the same.</li>
<li>I verified that my gpu is working correctly, I also was able to run the torchrl PPO Tutorial successfully on gpu.</li>
<li>i ran the scripts on my workplaces gpu server as well as on a separate windows laptop using torchrl-nightly and the behaviour still persists</li>
</ul>
<p>This leads me to conclude that I have some hidden bug in my script or did not use torchrl correctly. I would be really thankful for any pointer on what to look at.</p>
<p>Note: I also posted this on the torchrl github under discussion, but hope that it maybe finds more attention here.</p>
<pre><code class="lang-auto">import logging
import time

import torch.nn as nn
import torch.optim as optim

from tensordict.nn import TensorDictSequential as TDSeq

from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyTensorStorage, ReplayBuffer
from torchrl.envs import GymEnv, TransformedEnv, StepCounter, RewardSum, Compose
from torchrl.modules import EGreedyModule, QValueActor
from torchrl.objectives import DQNLoss, SoftUpdate


logging.basicConfig(level=logging.INFO)
my_logger = logging.getLogger(__name__)


ENV_NAME = "CartPole-v1"
DEVICE = "cuda"

INIT_RND_STEPS = 5_120
FRAMES_PER_BATCH = 128
BUFFER_SIZE = 100_000

GAMMA = 0.98
OPTIM_STEPS = 10
BATCH_SIZE = 128

SOFTU_EPS = 0.99
LR = 0.02


class Net(nn.Module):
    def __init__(self, obs_size: int, n_actions: int) -&gt; None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, 128),
            nn.ReLU(),
            nn.Linear(128, n_actions),
        )

    def forward(self, x):
        orig_shape_unbatched = len(x.shape) == 1
        if orig_shape_unbatched:
            x = x.unsqueeze(0)

        out = self.net(x)

        if orig_shape_unbatched:
            out = out.squeeze(0)
        return out


def make_env(env_name: str):
    return TransformedEnv(
        GymEnv(env_name),
        Compose(
            StepCounter(),
            RewardSum()
        )
    )


if __name__ == "__main__":
    env = make_env(ENV_NAME)
    n_obs = env.observation_spec["observation"].shape[-1]
    n_act = env.action_spec.shape[-1]

    net = Net(n_obs, n_act).to(device=DEVICE)
    agent = QValueActor(net, spec=env.action_spec)
    policy_explore = EGreedyModule(env.action_spec)
    agent_explore = TDSeq(agent, policy_explore)

    collector = SyncDataCollector(
        env,
        agent_explore,
        frames_per_batch=FRAMES_PER_BATCH,
        init_random_frames=INIT_RND_STEPS,
        device=DEVICE,
    )
    exp_buffer = ReplayBuffer(storage=LazyTensorStorage(BUFFER_SIZE))

    loss = DQNLoss(value_network=agent, action_space=env.action_spec, delay_value=True)
    loss.make_value_estimator(gamma=GAMMA)
    target_updater = SoftUpdate(loss, eps=SOFTU_EPS)
    optimizer = optim.Adam(loss.parameters(), lr=LR)

    total_count = 0
    total_episodes = 0
    t0 = time.time()
    for i, data in enumerate(collector):
        exp_buffer.extend(data)
        max_length = exp_buffer["next", "step_count"].max()
        max_reward = exp_buffer["next", "episode_reward"].max()
        if len(exp_buffer) &gt; INIT_RND_STEPS:
            for _ in range(OPTIM_STEPS):
                optimizer.zero_grad()
                sample = exp_buffer.sample(batch_size=BATCH_SIZE)
                sample = sample.to(DEVICE)

                loss_vals = loss(sample)
                loss_vals["loss"].backward()
                optimizer.step()

                agent_explore[1].step(data.numel())
                target_updater.step()

                total_count += data.numel()
                total_episodes += data["next", "done"].sum()

            if i % 10 == 0:
                my_logger.info(f"Step: {i}, max. count / epi reward: {max_length} / {max_reward}.")

        if max_length &gt; 200:
            t1 = time.time()
            my_logger.info(f"SOLVED in {t1 - t0}s!! MaxLen: {max_length}!")
            my_logger.info(f"With {max_reward} Reward!")
            my_logger.info(f"In {total_episodes} Episodes!")
            break
</code></pre>

username: vmoens
cooked: <p>There seems to be a bug with the data transfers in the collector, I’ll give it a shot next week and keep you posted!</p>

username: ozhicha
cooked: <p>Thank you, really appreciate it!</p>

username: vmoens
cooked: <p>Hello!<br>
I solved it in <a href="https://github.com/pytorch/rl/pull/2703" rel="noopener nofollow ugc">this pr</a><br>
The issue was that your policy wasn’t completely on CUDA. Part of it was still on CPU (namely the exploration module). Then when you ask the collector to run it, it sends everything to cuda. Here there was an issue which made the collector lose track of the origin tensor (this is what the PR is solving).<br>
You need to correct your script a bit though:</p>
<ul>
<li>you could do <code>agent_explore = agent_explore.to(device)</code>, in which case you don’t need the PR</li>
<li>If you don’t do that, use the PR (nightly build) and add a call to <code>collector.update_policy_weights_()</code> just after your model update</li>
</ul>
<pre data-code-wrap="python"><code class="lang-python">[...]
                total_count += data.numel()
                total_episodes += data["next", "done"].sum()

            if i % 10 == 0:
                my_logger.info(f"Step: {i}, max. count / epi reward: {max_length} / {max_reward}.")
        collector.update_policy_weights_()
</code></pre>
<p>That will copy the CPU buffers on GPU.</p>
<p>I also spotted a bug when you have partial devices (ie one for the policy and one for the env) which I fixed in <a href="https://github.com/pytorch/rl/pull/2704" class="inline-onebox" rel="noopener nofollow ugc">[BugFix] Fix device transfer for collectors with init_random_frames mixed devices by vmoens · Pull Request #2704 · pytorch/rl · GitHub</a></p>
<p>With these changes I solve the whole thing in a similar number of iterations in every case.</p>
<p>LMK if that works!</p>

username: ozhicha
cooked: <p>Thank you for looking into this issue!</p>
<ul>
<li>I can confirm that <code>agent_explore = agent_explore.to(device)</code> fixes the problem.</li>
</ul>
<p><strong>Edit: I tested on the lastest nightly not on the latest github changes thats why i had behaviour below. In actuality the proposed fix works!</strong></p>
<ul>
<li>adding <code>collector.update_policy_weights_()</code> fixes the issue only if device in collector is set to cpu, as in below</li>
</ul>
<pre><code class="lang-auto">collector = SyncDataCollector(
        env,
        agent_explore,
        frames_per_batch=FRAMES_PER_BATCH,
        init_random_frames=INIT_RND_STEPS,
        device="cpu",
    )
</code></pre>
<p>if  <code>device="cuda"</code> training still doesn’t converge.</p>
<p>For me this is fine and I would have been happy with just pushing the policy on to the gpu as fix, but maybe one would want to enforce that if device in collector is set, than policy has to be on the same device.</p>
<p>Again, thank you for looking into this! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: vmoens
cooked: <aside class="quote no-group" data-username="ozhicha" data-post="5" data-topic="215547">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ozhicha/48/75169_2.png" class="avatar"> ozhicha:</div>
<blockquote>
<p>if <code>device="cuda"</code> training still doesn’t converge.</p>
</blockquote>
</aside>
<p>it will only if you pull the latest changes from github!</p>

username: ozhicha
cooked: <p>my bad you are right! I tested on the nightly version of torchrl but not with latest changes from github.</p>
<p>Can confirm, it work!</p>