[rllib] Custom environment observations always have the dtype float32

[internetcoffeephone]

ray[rllib] 0.8.3
Python 3.6
tensorflow-gpu 2.1.0

What is the problem?

When using a custom environment with an observation with a dtype other than float32, e.g. uint8, the observation type is changed to float32. This causes models that expect other data types to fail.

This happens due to a hardcoded float32 in dynamic_tf_policy.py.

When using a Dict observation space (a set of simpler gym spaces), it also happens in preprocessors.py.

Reproduction

import numpy as np
import ray
from gym.spaces import Box, Dict, Discrete
from ray import tune
from ray.rllib.env import BaseEnv
from ray.rllib.utils import try_import_tf
from ray.tune.registry import register_env

tf = try_import_tf()


class TestEnv(BaseEnv):
    def poll(self):
        return {}, [], [], [], []

    def send_actions(self, action_dict):
        pass

    observation_space = Dict({"curr_obs": Box(low=0, high=255, shape=(1, 1), dtype=np.uint8)})
    action_space = Discrete(1)


def env_creator(_):
    return TestEnv()


env_name = "test_env"
register_env(env_name, env_creator)

ray.init(local_mode=True)
tune.run("PPO", config={"env": env_name})

Note: You will have to set the breakpoints in the affected lines yourself. I wanted to provide a simple custom model that takes in only uint8, but couldn't get it to work.

[internetcoffeephone]

It seems that incorrect sample dtypes are used in other places as well:

• ModelCatalog Solved by Change get_action_shape to avoid unexpected behavior #8424
• TupleFlatteningPreprocessor
• DictFlatteningPreprocessor
• Deterministic action sampling
• Action sampling
• Sample batch builder

[stale]

Hi, I'm a bot from the Ray team :)

To help human contributors to focus on more relevant issues, I will automatically add the stale label to issues that have had no activity for more than 4 months.

If there is no further activity in the 14 days, the issue will be closed!

• If you'd like to keep the issue open, just leave any comment, and the stale label will be removed!
• If you'd like to get more attention to the issue, please tag one of Ray's contributors.

You can always ask for help on our discussion forum or Ray's public slack channel.

[internetcoffeephone]

Comment to remove stale label.

[stale]

Hi, I'm a bot from the Ray team :)

To help human contributors to focus on more relevant issues, I will automatically add the stale label to issues that have had no activity for more than 4 months.

If there is no further activity in the 14 days, the issue will be closed!

• If you'd like to keep the issue open, just leave any comment, and the stale label will be removed!
• If you'd like to get more attention to the issue, please tag one of Ray's contributors.

You can always ask for help on our discussion forum or Ray's public slack channel.

[internetcoffeephone]

This comment is to remove the stale label.

[stale]

Hi, I'm a bot from the Ray team :)

To help human contributors to focus on more relevant issues, I will automatically add the stale label to issues that have had no activity for more than 4 months.

If there is no further activity in the 14 days, the issue will be closed!

• If you'd like to keep the issue open, just leave any comment, and the stale label will be removed!
• If you'd like to get more attention to the issue, please tag one of Ray's contributors.

You can always ask for help on our discussion forum or Ray's public slack channel.

[internetcoffeephone]

This comment is to remove the stale label.

[avnishn]

closing because after running the repro script, it looks like we chased this bug away.