Training converges on device="cpu" but never on device="cuda" #2691

ozhicha · 2025-01-12T07:27:18Z

ozhicha
Jan 12, 2025

Hi,
for the script below my training always converges on CPU and solves the environment. However if I use my GPU I can't see any progress. I'm basically following the tutorials and at this point have no idea where my bug could be. I'm also pretty sure my GPU works as an vanilla implementation without torchrl for a dqn on pong worked without problem.
For my python env: gym==0.26.2, torch==2.5.1, torchrl==0.5.0 with cuda 12.4

Script:

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 QValueActor, EGreedyModule
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) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, 128),
            nn.ReLU(),
            nn.Linear(128, n_actions),
        )

    def forward(self, x):
        return self.net(x)


def make_env(env_name: str):
    return TransformedEnv(
        GymEnv(env_name, device=DEVICE),
        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]

    module = Net(n_obs, n_act).to(device=DEVICE)
    agent = QValueActor(module, 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,
        split_trajs=False,
        device=DEVICE,
        storing_device=DEVICE,
        policy_device=DEVICE,
    )
    exp_buffer = ReplayBuffer(storage=LazyTensorStorage(BUFFER_SIZE))

    loss = DQNLoss(value_network=agent, 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) > 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:
                my_logger.info(f"Step: {i}, max. count / epi reward: {max_length} / {max_reward}.")

        if max_length > 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

Output for DEVICE="cuda":
INFO:main:Step: 41, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 42, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 43, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 44, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 45, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 46, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 47, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 48, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 49, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 51, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 52, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 53, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 54, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 55, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 56, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 57, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 58, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 59, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 61, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 62, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 63, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 64, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 65, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 66, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 67, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 68, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 69, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 71, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 72, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 73, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 74, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 75, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 76, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 77, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 78, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 79, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 81, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 82, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 83, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 84, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 85, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 86, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 87, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 88, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 89, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 91, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 92, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 93, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 94, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 95, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 96, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 97, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 98, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 99, max. count / epi reward: 72 / 72.0.
INFO:main:Step: 101, max. count / epi reward: 72 / 72.0.

Output for DEVICE="cpu":
INFO:main:Step: 41, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 42, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 43, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 44, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 45, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 46, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 47, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 48, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 49, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 51, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 52, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 53, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 54, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 55, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 56, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 57, max. count / epi reward: 84 / 84.0.
INFO:main:Step: 58, max. count / epi reward: 93 / 93.0.
INFO:main:Step: 59, max. count / epi reward: 93 / 93.0.
INFO:main:Step: 61, max. count / epi reward: 93 / 93.0.
INFO:main:Step: 62, max. count / epi reward: 93 / 93.0.
INFO:main:Step: 63, max. count / epi reward: 93 / 93.0.
INFO:main:Step: 64, max. count / epi reward: 93 / 93.0.
INFO:main:Step: 65, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 66, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 67, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 68, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 69, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 71, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 72, max. count / epi reward: 110 / 110.0.
INFO:main:Step: 73, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 74, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 75, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 76, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 77, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 78, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 79, max. count / epi reward: 122 / 122.0.
INFO:main:Step: 81, max. count / epi reward: 129 / 129.0.
INFO:main:Step: 82, max. count / epi reward: 145 / 145.0.
INFO:main:Step: 83, max. count / epi reward: 145 / 145.0.
INFO:main:Step: 84, max. count / epi reward: 145 / 145.0.
INFO:main:Step: 85, max. count / epi reward: 145 / 145.0.
INFO:main:Step: 86, max. count / epi reward: 145 / 145.0.
INFO:main:Step: 87, max. count / epi reward: 151 / 151.0.
INFO:main:Step: 88, max. count / epi reward: 173 / 173.0.
INFO:main:Step: 89, max. count / epi reward: 173 / 173.0.
INFO:main:Step: 91, max. count / epi reward: 173 / 173.0.
INFO:main:Step: 92, max. count / epi reward: 173 / 173.0.
INFO:main:Step: 93, max. count / epi reward: 173 / 173.0.
INFO:main:Step: 94, max. count / epi reward: 225 / 225.0.
INFO:main:SOLVED in 4.681804656982422s!! MaxLen: 225!
INFO:main:With 225.0 Reward!
INFO:main:In 3620 Episodes!

Answered by ozhicha

Jan 20, 2025

Answer can be found here: https://discuss.pytorch.org/t/training-converges-on-cpu-but-never-on-gpu/215547/4

View full answer

vmoens · 2025-01-13T14:41:57Z

vmoens
Jan 13, 2025
Collaborator

Hello! There could be some issues with H2D transfer on older versions of torchrl and tensordict, can you try with something > v0.6.0?

1 reply

ozhicha Jan 13, 2025
Author

Thanks, I did rerun the code on torchrl 0.6.0 with tensordict 0.6.0 but unfortunately I still have the same behaviour :(

Edit:
I have retried with tensordict 0.6.2 as well as with the current torchrl-nightly version, however I'm still getting the same behaviour. I also did a rerun of my dqn implementation without torchrl and it is converging. I also used copied the dqn trainer example from the tutorial section, which returned matching values to the tutorial I guess (see below). So it looks like there is some bug in my implementation on top.

r_training: 0.4295, rewards: 0.1000, total_rewards: 6.6667: : 5024it [00:51, 112.12it/s]File: /tmp/nix-shell-2288-0/tmp7887nyfr/dqn_exp_bb9d3f2e-d24e-11ef-941e-74563cbfccf9/scalars/r_training.csv
512,0.42945271730422974
1024,0.40816524624824524
1536,0.42945271730422974
2048,0.4597761034965515
2560,0.4597761034965515
3072,0.42945271730422974
3584,0.4597761034965515
4096,0.42945271730422974
4608,0.4597761034965515

File: /tmp/nix-shell-2288-0/tmp7887nyfr/dqn_exp_bb9d3f2e-d24e-11ef-941e-74563cbfccf9/scalars/total_rewards.csv
3232,6.666666507720947

File: /tmp/nix-shell-2288-0/tmp7887nyfr/dqn_exp_bb9d3f2e-d24e-11ef-941e-74563cbfccf9/scalars/grad_norm_0.csv
512,3.198786497116089
1024,2.0401320457458496
1536,1.9391694068908691
2048,3.432772159576416
2560,2.8411054611206055
3072,3.8100593090057373
3584,3.458214521408081
4096,4.958317279815674
4608,3.726311206817627

File: /tmp/nix-shell-2288-0/tmp7887nyfr/dqn_exp_bb9d3f2e-d24e-11ef-941e-74563cbfccf9/scalars/rewards.csv
3232,0.10000000894069672

File: /tmp/nix-shell-2288-0/tmp7887nyfr/dqn_exp_bb9d3f2e-d24e-11ef-941e-74563cbfccf9/scalars/optim_steps.csv
512,128.0
1024,256.0
1536,384.0
2048,512.0
2560,640.0
3072,768.0
3584,896.0
4096,1024.0
4608,1152.0

File: /tmp/nix-shell-2288-0/tmp7887nyfr/dqn_exp_bb9d3f2e-d24e-11ef-941e-74563cbfccf9/scalars/loss.csv
512,0.2476823776960373
1024,0.17451828718185425
1536,0.21816502511501312
2048,0.3180880844593048
2560,0.3251710534095764
3072,0.36744266748428345
3584,0.2928861975669861
4096,0.2788560688495636
4608,0.3400101661682129

ozhicha · 2025-01-20T13:35:05Z

ozhicha
Jan 20, 2025
Author

Answer can be found here: https://discuss.pytorch.org/t/training-converges-on-cpu-but-never-on-gpu/215547/4

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Training converges on device="cpu" but never on device="cuda" #2691

{{title}}

Replies: 2 comments 1 reply

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

Select a reply

Training converges on device="cpu" but never on device="cuda" #2691

ozhicha Jan 12, 2025

Replies: 2 comments · 1 reply

vmoens Jan 13, 2025 Collaborator

ozhicha Jan 13, 2025 Author

ozhicha Jan 20, 2025 Author

ozhicha
Jan 12, 2025

Replies: 2 comments 1 reply

vmoens
Jan 13, 2025
Collaborator

ozhicha Jan 13, 2025
Author

ozhicha
Jan 20, 2025
Author