# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2
# 
# Path: android_env/wrappers/base_wrapper.py
# class BaseWrapper(dm_env.Environment):
#   def __init__(self, env):
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, action: Any) -> dm_env.TimeStep:
#   def _reset_state(self):
#   def _process_action(self, action: Any) -> Any:
#   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
#   def observation_spec(self) -> Dict[str, specs.Array]:
#   def action_spec(self) -> Dict[str, specs.Array]:
#   def _wrapper_logs(self) -> Dict[str, Any]:
#   def android_logs(self) -> Dict[str, Any]:
#   def raw_env(self):
#   def __getattr__(self, attr):
#   def close(self):


from typing import Optional, Sequence, Dict
from android_env.components import action_type
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np
                     dtype=self._parent_action_spec['touch_position'].dtype)
    }

  def _get_action_type(self, action_id: int) -> action_type.ActionType:
    """Compute action type corresponding to the given action_id.

    When `self._redundant_actions` == True the `grid_size` is "broadcast" over
    all the possible actions so you end up with `grid_size` discrete actions
    of type 0, `grid_size` discrete actions of type 1, etc. for all action
    types.

    When `self._redundant_actions` == False the first `grid_size` actions are
    reserved for "touch" and the rest are just added (NOT multiplied) to the
    total number of discrete actions (exactly one of LIFT and REPEAT).

    Args:
      action_id: A discrete action.
    Returns:
      action_type: The action_type of the action.
    """

    if self._redundant_actions:
      assert action_id < self._num_action_types * self._grid_size
      return action_id // self._grid_size

    else:
      assert action_id <= self._grid_size + 1
      if action_id < self._grid_size:
        return action_type.ActionType.TOUCH
      elif action_id == self._grid_size:

# Path: android_env/components/action_type.py
# class ActionType(enum.IntEnum):
#   TOUCH = 0
#   LIFT = 1
#   REPEAT = 2
# 
# Path: android_env/wrappers/base_wrapper.py
# class BaseWrapper(dm_env.Environment):
#   def __init__(self, env):
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, action: Any) -> dm_env.TimeStep:
#   def _reset_state(self):
#   def _process_action(self, action: Any) -> Any:
#   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
#   def observation_spec(self) -> Dict[str, specs.Array]:
#   def action_spec(self) -> Dict[str, specs.Array]:
#   def _wrapper_logs(self) -> Dict[str, Any]:
#   def android_logs(self) -> Dict[str, Any]:
#   def raw_env(self):
#   def __getattr__(self, attr):
#   def close(self):


# Path: android_env/wrappers/discrete_action_wrapper.py
from typing import Optional, Sequence, Dict
from android_env.components import action_type
from android_env.wrappers import base_wrapper
from dm_env import specs
import dm_env
import numpy as np

                     dtype=self._parent_action_spec['touch_position'].dtype)
    }

  def _get_action_type(self, action_id: int) -> action_type.ActionType:
    """Compute action type corresponding to the given action_id.

    When `self._redundant_actions` == True the `grid_size` is "broadcast" over
    all the possible actions so you end up with `grid_size` discrete actions
    of type 0, `grid_size` discrete actions of type 1, etc. for all action
    types.

    When `self._redundant_actions` == False the first `grid_size` actions are
    reserved for "touch" and the rest are just added (NOT multiplied) to the
    total number of discrete actions (exactly one of LIFT and REPEAT).

    Args:
      action_id: A discrete action.
    Returns:
      action_type: The action_type of the action.
    """

    if self._redundant_actions:
      assert action_id < self._num_action_types * self._grid_size
      return action_id // self._grid_size

    else:
      assert action_id <= self._grid_size + 1
      if action_id < self._grid_size:
        return action_type.ActionType.TOUCH
      elif action_id == self._grid_size:

        return action_type.ActionType.LIFT