class_docstr: Stationary Stochastic Bandit environment with per-arm features.
class_docstr_tok: ["Stationary","Stochastic","Bandit","environment","with","per-arm","features","."]
docstr: We need to override this function because the reward dtype can be int.
docstr_tok: ["We","need","to","override","this","function","because","the","reward","dtype","can","be","int","."]
code: def _step(self, action):
      """We need to override this function because the reward dtype can be int."""
      # TODO(b/199824775): The trajectory module assumes all reward is float32.
      # Sort this out with TF-Agents.
      output = super(RankingPyEnvironment, self)._step(action)
      reward = output.reward
      new_reward = tf.nest.map_structure(
          lambda x, t: x.astype(t), reward, self.reward_spec()
      )
      return ts.TimeStep(
          step_type=output.step_type,
          reward=new_reward,
          discount=output.discount,
          observation=output.observation,
      )
  
imports: from typing import Callable, Optional, Sequence, Text
  import numpy
  import tensorflow
  from tf_agents.bandits.environments import bandit_py_environment
  from tf_agents.bandits.specs import utils
  from tf_agents.specs import array_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
label_desc: Use image node_id 7 for calling the RankingPyEnvironment obj's underlying member method code with example usage: obj._step(action) and returns: ts