<commit_before># Copyright 2020 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A simple human player for testing a Melting Pot level.

Use `WASD` keys to move the character around.
Use `Q and E` to turn the character.
Use `SPACE` to fire the zapper.
Use `TAB` to switch between players.
"""

import collections
import enum

from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple
import dm_env

from ml_collections import config_dict
import numpy as np
import pygame

import dmlab2d
from meltingpot.python.utils.substrates import builder

WHITE = (255, 255, 255)

MOVEMENT_MAP = {
    'NONE': 0,
    'FORWARD': 1,
    'RIGHT': 2,
    'BACKWARD': 3,
    'LEFT': 4,
}


EnvBuilder = Callable[..., dmlab2d.Environment]  # Only supporting kwargs.

ActionMap = Mapping[str, Callable[[], int]]


class RenderType(enum.Enum):
  NONE = 0
  PYGAME = 1


def get_random_direction() -> int:
  """Gets a random direction."""
  return np.random.choice(list(MOVEMENT_MAP.values()))


def get_random_turn() -> int:
  """Gets a random turn."""
  return np.random.choice([-1, 0, 1])


def get_random_fire() -> int:
  """Gets a random fire."""
  return np.random.choice([0, 1])


def get_direction_pressed() -> int:
  """Gets direction pressed."""
  key_pressed = pygame.key.get_pressed()
  if key_pressed[pygame.K_UP] or key_pressed[pygame.K_w]:
    return MOVEMENT_MAP['FORWARD']
  if key_pressed[pygame.K_RIGHT] or key_pressed[pygame.K_d]:
    return MOVEMENT_MAP['RIGHT']
  if key_pressed[pygame.K_DOWN] or key_pressed[pygame.K_s]:
    return MOVEMENT_MAP['BACKWARD']
  if key_pressed[pygame.K_LEFT] or key_pressed[pygame.K_a]:
    return MOVEMENT_MAP['LEFT']
  return MOVEMENT_MAP['NONE']


def get_turn_pressed() -> int:
  """Calculates turn increment."""
  key_pressed = pygame.key.get_pressed()
  if key_pressed[pygame.K_DELETE] or key_pressed[pygame.K_q]:
    return -1
  if key_pressed[pygame.K_PAGEDOWN] or key_pressed[pygame.K_e]:
    return 1
  return 0


def get_space_key_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_SPACE] else 0


def get_key_number_pressed() -> int:
  number_keys = [pygame.K_0, pygame.K_1, pygame.K_2, pygame.K_3, pygame.K_4,
                 pygame.K_5, pygame.K_6, pygame.K_7, pygame.K_8, pygame.K_9]
  for num in range(len(number_keys)):
    if pygame.key.get_pressed()[number_keys[num]]:
      return num
  return -1


def get_key_number_one_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_1] else 0


def get_key_number_two_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_2] else 0


def get_key_number_three_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_3] else 0


def get_key_number_four_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_4] else 0


def get_key_number_five_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_5] else 0


def get_left_control_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_LCTRL] else 0


def get_left_shift_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_LSHIFT] else 0


def get_right_shift_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_RSHIFT] else 0


def get_key_c_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_c] else 0


def get_key_z_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_z] else 0


def get_key_x_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_x] else 0


def _split_key(key: str) -> Tuple[str, str]:
  """Splits the key into player index and name."""
  return tuple(key.split('.', maxsplit=1))


def _get_rewards(timestep: dm_env.TimeStep) -> Mapping[str, float]:
  """Gets the list of rewards, one for each player."""
  rewards = {}
  for key in timestep.observation.keys():
    if key.endswith('.REWARD'):
      player_prefix, name = _split_key(key)
      if name == 'REWARD':
        rewards[player_prefix] = timestep.observation[key]
  return rewards


class ActionReader(object):
  """Convert keyboard actions to environment actions."""

  def __init__(self, env: dmlab2d.Environment, action_map: ActionMap):
    # Actions are named "<player_prefix>.<action_name>"
    self._action_map = action_map
    self._action_spec = env.action_spec()
    assert isinstance(self._action_spec, dict)
    self._action_names = set()
    for action_key in self._action_spec.keys():
      _, action_name = _split_key(action_key)
      self._action_names.add(action_name)

  def step(self, player_prefix: str) -> Sequence[Sequence[Mapping[str, int]]]:
    """Update the actions of player `player_prefix`."""
    actions = {action_key: 0 for action_key in self._action_spec.keys()}
    for action_name in self._action_names:
      actions[f'{player_prefix}.{action_name}'] = self._action_map[
          action_name]()
    return actions

<commit_msg>Update incorrect type annotation.<commit_after># Copyright 2020 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A simple human player for testing a Melting Pot level.

Use `WASD` keys to move the character around.
Use `Q and E` to turn the character.
Use `SPACE` to fire the zapper.
Use `TAB` to switch between players.
"""

import collections
import enum

from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple
import dm_env

from ml_collections import config_dict
import numpy as np
import pygame

import dmlab2d
from meltingpot.python.utils.substrates import builder

WHITE = (255, 255, 255)

MOVEMENT_MAP = {
    'NONE': 0,
    'FORWARD': 1,
    'RIGHT': 2,
    'BACKWARD': 3,
    'LEFT': 4,
}


EnvBuilder = Callable[..., dmlab2d.Environment]  # Only supporting kwargs.

ActionMap = Mapping[str, Callable[[], int]]


class RenderType(enum.Enum):
  NONE = 0
  PYGAME = 1


def get_random_direction() -> int:
  """Gets a random direction."""
  return np.random.choice(list(MOVEMENT_MAP.values()))


def get_random_turn() -> int:
  """Gets a random turn."""
  return np.random.choice([-1, 0, 1])


def get_random_fire() -> int:
  """Gets a random fire."""
  return np.random.choice([0, 1])


def get_direction_pressed() -> int:
  """Gets direction pressed."""
  key_pressed = pygame.key.get_pressed()
  if key_pressed[pygame.K_UP] or key_pressed[pygame.K_w]:
    return MOVEMENT_MAP['FORWARD']
  if key_pressed[pygame.K_RIGHT] or key_pressed[pygame.K_d]:
    return MOVEMENT_MAP['RIGHT']
  if key_pressed[pygame.K_DOWN] or key_pressed[pygame.K_s]:
    return MOVEMENT_MAP['BACKWARD']
  if key_pressed[pygame.K_LEFT] or key_pressed[pygame.K_a]:
    return MOVEMENT_MAP['LEFT']
  return MOVEMENT_MAP['NONE']


def get_turn_pressed() -> int:
  """Calculates turn increment."""
  key_pressed = pygame.key.get_pressed()
  if key_pressed[pygame.K_DELETE] or key_pressed[pygame.K_q]:
    return -1
  if key_pressed[pygame.K_PAGEDOWN] or key_pressed[pygame.K_e]:
    return 1
  return 0


def get_space_key_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_SPACE] else 0


def get_key_number_pressed() -> int:
  number_keys = [pygame.K_0, pygame.K_1, pygame.K_2, pygame.K_3, pygame.K_4,
                 pygame.K_5, pygame.K_6, pygame.K_7, pygame.K_8, pygame.K_9]
  for num in range(len(number_keys)):
    if pygame.key.get_pressed()[number_keys[num]]:
      return num
  return -1


def get_key_number_one_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_1] else 0


def get_key_number_two_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_2] else 0


def get_key_number_three_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_3] else 0


def get_key_number_four_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_4] else 0


def get_key_number_five_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_5] else 0


def get_left_control_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_LCTRL] else 0


def get_left_shift_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_LSHIFT] else 0


def get_right_shift_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_RSHIFT] else 0


def get_key_c_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_c] else 0


def get_key_z_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_z] else 0


def get_key_x_pressed() -> int:
  return 1 if pygame.key.get_pressed()[pygame.K_x] else 0


def _split_key(key: str) -> Tuple[str, str]:
  """Splits the key into player index and name."""
  return tuple(key.split('.', maxsplit=1))


def _get_rewards(timestep: dm_env.TimeStep) -> Mapping[str, float]:
  """Gets the list of rewards, one for each player."""
  rewards = {}
  for key in timestep.observation.keys():
    if key.endswith('.REWARD'):
      player_prefix, name = _split_key(key)
      if name == 'REWARD':
        rewards[player_prefix] = timestep.observation[key]
  return rewards


class ActionReader(object):
  """Convert keyboard actions to environment actions."""

  def __init__(self, env: dmlab2d.Environment, action_map: ActionMap):
    # Actions are named "<player_prefix>.<action_name>"
    self._action_map = action_map
    self._action_spec = env.action_spec()
    assert isinstance(self._action_spec, dict)
    self._action_names = set()
    for action_key in self._action_spec.keys():
      _, action_name = _split_key(action_key)
      self._action_names.add(action_name)

  def step(self, player_prefix: str) -> Mapping[str, int]:
    """Update the actions of player `player_prefix`."""
    actions = {action_key: 0 for action_key in self._action_spec.keys()}
    for action_name in self._action_names:
      actions[f'{player_prefix}.{action_name}'] = self._action_map[
          action_name]()
    return actions


{'fraction_non_alphanumeric': 0.11883116883116883, 'fraction_numerical': 0.014285714285714285, 'mean_word_length': 4.1757081132981275, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 49, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6884817', 'n_tokens_mistral': 3866, 'n_tokens_neox': 3688, 'n_words': 1027}