Why does a dtype mismatch occur in one-hot encoding environments?

[liangping]

Hello everyone,

I've encountered a dtype mismatch issue with env.rollout or SyncDataCollector when using a one-hot encoding environment like CliffWalking. However, the code works as expected with other environments, such as 'CartPole-v1'.

Am I missing something?

env = GymEnv("CliffWalking-v0")

env = TransformedEnv(env,  StepCounter(max_steps=100))

n_obs = env.observation_spec["observation"].shape[-1]
n_act = env.action_spec.shape[-1]
actor = Actor(MLP(in_features=n_obs, out_features=n_act, num_cells=[32, 32]))

env.rollout(max_steps=1, policy=actor)

  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torchrl/modules/models/models.py", line 300, in forward
    out = super().forward(*inputs)
  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torch/nn/modules/container.py", line 240, in forward
    input = module(input)
  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/ping/.venv/hello/lib/python3.13/site-packages/torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
           ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: mat1 and mat2 must have the same dtype, but got Long and Float

[liangping]

This issue should be solved by adding a transform.

env = TransformedEnv(GymEnv(env_name), Compose(DTypeCastTransform(dtype_in=torch.int64, dtype_out=torch.float32, in_keys=["observation"]), StepCounter()))