fix bug of road_traffic senario when using cuda

[Schopenhauer-loves-Hegel]

hello, I discovered that the road traffic scenario throws an error when running on CUDA due to inconsistent device usage in road_traffic.py, so I modified the file to fix this bug.

The bug can be reproduced with the following code:

import vmas  
import time  

def run_env_steps(device="cuda", num_steps=100):  
    env = vmas.make_env(  
        scenario="road_traffic",  
        num_envs=4,   
        device=device,  
        continuous_actions=True,  
        wrapper=None,  
        max_steps=None,  
        seed=None,  
        dict_spaces=False,  
        grad_enabled=False,  
        terminated_truncated=False,  
    )  

    start_time = time.time()  

    obs = env.reset()  
    for i in range(num_steps):  
        random_actions = env.get_random_actions()  
        obs, rewards, dones, info = env.step(random_actions)  

        if dones.any():  
            obs = env.reset()  

    end_time = time.time()  
    total_time = end_time - start_time  

    print(f"\nRunning {num_steps} steps on {device}:")  
    print(f"Total time: {total_time:.4f} seconds")  
    print(f"Average time per step: {(total_time/num_steps)*1000:.4f} ms")  

    return total_time  

# Test CPU vs. GPU performance  
print("Starting performance test...")  
cpu_time = run_env_steps(device="cpu", num_steps=100)  
gpu_time = run_env_steps(device="cuda", num_steps=100)  

print("\nPerformance comparison:")  
print(f"CPU total time: {cpu_time:.4f} seconds")  
print(f"GPU total time: {gpu_time:.4f} seconds")  
print(f"GPU speedup over CPU: {cpu_time/gpu_time:.2f}x") 

[matteobettini]

Thanks a mil for this!

[matteobettini]

cc @Jianye-Xu