$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 89. Improved: $\large\color{#35bf28}14$. Worsened: $\large\color{#d91a1a}7$.

does this make sense to use with a parent environment? when is this transform preferred to env.to(device)?

This is to address #1198 where the issue is that if the env naturally sits on MPS we can't use float64. So first you must transform the data into float32 and then cast to device. Doing env.to(device) will not work but this transform will.

does this make sense to use with a parent environment? when is this transform preferred to env.to(device)?

This is to address #1198 where the issue is that if the env naturally sits on MPS we can't use float64. So first you must transform the data into float32 and then cast to device. Doing env.to(device) will not work but this transform will.

Could you please explain exactly how to do this? I am confused. I am following this tutorial https://pytorch.org/tutorials/intermediate/reinforcement_ppo.html and get the MPS flaot64 error when running line: base_env = GymEnv("InvertedDoublePendulum-v4", device=device, frame_skip=frame_skip) What code exactly shall I write to correct this error please?

Can I ask what the value of device is in your case?

device="mps".

I believe I solved it with this (after quite a few hours of trying different things!!!):

base_env = GymEnv("InvertedDoublePendulum-v4", device="cpu", frame_skip=frame_skip)
env = TransformedEnv(
    base_env,
    Compose(
        ObservationNorm(in_keys=["observation"]), # normalise observations (make it about Standard Normal)
        DoubleToFloat(),   
        StepCounter(),                            # count the number of steps before the environment is terminated
        DeviceCastTransform(device=device, orig_device="cpu"),
    ),
)
print(env.device) # gives mps now

Could you please kindly confirm if what Ive done is correct? I am on Apple M2 max trying to use MPS.

Everything was progressing smoothly through the tutorial: https://pytorch.org/tutorials/intermediate/reinforcement_ppo.html however at this code I get an error again about MPS. Please could you kindly advise syntax to solve this issue? I double checked everything seems to be on mps, so I don't understand where the error is coming from.

collector = SyncDataCollector(
    env,
    policy_module,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device,
)


  File ~/anaconda3/envs/gpu-torch-rl/lib/python3.10/site-packages/torch/nn/modules/module.py:1143 in convert
    return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.


env.device
Out[23]: device(type='mps')

@EkaterinaAbramova sorry you had this terrible experience, we should document things better for MPS.
It's something we're actively looking at and any feedback is very welcome.
The way you implemented your env looks great to me!

Regarding the error I will have a look into it, shouldn't be too difficult to solve. Have you tried moving the ObservationNorm and StepCounter after the device casting transform?
like:

env = TransformedEnv(
    base_env,
    Compose(
        DoubleToFloat(),   
        DeviceCastTransform(device=device, orig_device="cpu"),
        ObservationNorm(in_keys=["observation"]), # normalise observations (make it about Standard Normal)
        StepCounter(),                            # count the number of steps before the environment is terminated
    ),
)

Like this the buffers in the ObservationNorm transform will sit on mps but with float32 and not float64.

@vmoens thank you for swiftly helping me, this issue is quite urgent, so very glad that you had the suggestion. It makes sense and I tried it, however this way around I get this error AttributeError: 'DoubleToFloat' object has no attribute 'init_stats' (it seems I need to pass some arguments, what shall I pass to be able to follow the tutorial please https://pytorch.org/tutorials/intermediate/reinforcement_ppo.html):

You need to call init_stats on the obs norm transform

env.transform[-2].init_stats(...)

Because the transform has changed place

OK I get it. So now that Ive indexed the correct location, the MPS issue is back.
The code you suggested doesn't work for me:

env = TransformedEnv(
    base_env,
    Compose(
        DoubleToFloat(),   
        DeviceCastTransform(device=device, orig_device="cpu"),
        ObservationNorm(in_keys=["observation"]), # normalise observations (make it about Standard Normal)
        StepCounter(),                            # count the number of steps before the environment is terminated
    ),
)
print(env.device)
env.transform[-2].init_stats(num_iter=1000, reduce_dim=0, cat_dim=0) 
I get the error: Cannot convert a MPS Tensor

TO RECAP: The original code was:

env = TransformedEnv(
    base_env,
    Compose(
        ObservationNorm(in_keys=["observation"]), # normalise observations (make it about Standard Normal)
        DoubleToFloat(),   
        StepCounter(),                            # count the number of steps before the environment is terminated
        DeviceCastTransform(device=device, orig_device="cpu"),
    ),
)
env.transform[0].init_stats(num_iter=1000, reduce_dim=0, cat_dim=0) 

This code runs, however I have issues further below in the tutorial where I can't run:

collector = SyncDataCollector(
    env,
    policy_module,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device,
)
Cannot convert a MPS Tensor

Could you please provide a solution? I am really in a rush now, Ive been trying to solve this issue for days. The tutorial I am working on is online, so if it helps maybe you could try the suggestions to make sure they resolve the issue? Thank you very much for your help, I really need to move past this ASAP please.

Ok so that's an interesting bug, which basically boils down to some internal machinery within rollout, resets and transforms.
To quickly unblock you: can you compute the stats manually with your env?

simple_env = TransformedEnv(
    base_env,
    Compose(
        DoubleToFloat(),   
        DeviceCastTransform(device=device, orig_device="cpu"),
   )
)
td0 = simple_env.rollout(100)
loc = td0["observation"].mean(dim=0)
scale = td0["observation"].std(dim=0)
env = TransformedEnv(
    base_env,
    Compose(
        DoubleToFloat(),   
        DeviceCastTransform(device=device, orig_device="cpu"),
        ObservationNorm(in_keys=["observation"], loc=loc, scale=scale),
        StepCounter(),
    ),
)

Hopefully that should help!
I should be able to put my hands on an apple silicon computer tomorrow morning if you're still stuck!

The suggestion didn't work Im afraid. Any other thing you could propose at this stage or only tomorrow?

So why in the version I provided above, everything is fine until I get down to the SyncDataCollector? Why is it failing there? What has not been yet converted to float32?

#1589 will solve your problem!

I have read through that thread, but I don't understand what I am supposed to do. Sorry!! Should I download the packages again from fresh? Or download something particular from that page?

Just wait until we merge it and then you can reinstall from got and things should work ok

@EkaterinaAbramova you should be good to go now!

@vmoens hey I started from scratch but things are not working. I made a post here but nobody replied yet please: #1198

I can reproduce your issue, let me push a fix!