<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.
"""Scenario factory."""

from typing import Collection, Mapping

from ml_collections import config_dict

from meltingpot.python import bot as bot_factory
from meltingpot.python import substrate as substrate_factory
from meltingpot.python.configs import scenarios as scenario_config
from meltingpot.python.utils.scenarios import population
from meltingpot.python.utils.scenarios import scenario as scenario_lib
from meltingpot.python.utils.scenarios import substrate_transforms

AVAILABLE_SCENARIOS = frozenset(scenario_config.SCENARIO_CONFIGS)

SCENARIOS_BY_SUBSTRATE: Mapping[
    str, Collection[str]] = scenario_config.scenarios_by_substrate(
        scenario_config.SCENARIO_CONFIGS)

PERMITTED_OBSERVATIONS = frozenset({
    'INVENTORY',
    'READY_TO_SHOOT',
    'RGB',
})

# TODO(b/227143834): Remove aliases once internal deps have been removed.
Scenario = scenario_lib.Scenario
PopulationObservables = population.PopulationObservables
ScenarioObservables = scenario_lib.ScenarioObservables
Population = population.Population


def get_config(scenario_name: str) -> config_dict.ConfigDict:
  """Returns a config for the specified scenario.

  Args:
    scenario_name: Name of the scenario. Must be in AVAILABLE_SCENARIOS.
  """
  if scenario_name not in AVAILABLE_SCENARIOS:
    raise ValueError(f'Unknown scenario {scenario_name!r}')
  scenario = scenario_config.SCENARIO_CONFIGS[scenario_name]
  substrate = substrate_factory.get_config(scenario.substrate)
  bots = {name: bot_factory.get_config(name) for name in scenario.bots}
  config = config_dict.create(
      substrate=substrate,
      bots=bots,
      is_focal=scenario.is_focal,
      num_players=sum(scenario.is_focal),
      num_bots=len(scenario.is_focal) - sum(scenario.is_focal),
  )
  return config.lock()


def build(config: config_dict.ConfigDict) -> scenario_lib.Scenario:
  """Builds a scenario for the given config.

  Args:
    config: config resulting from `get_config`.

  Returns:
    The test scenario.
  """
  substrate = substrate_factory.build(config.substrate)

  # Add observations needed by some bots. These are removed for focal players.
  permitted_observations = (
      set(substrate.observation_spec()[0]) & PERMITTED_OBSERVATIONS)
  substrate = substrate_transforms.with_tf1_bot_required_observations(substrate)

  background_population = population.Population(
      policies={
          bot_name: bot_factory.build(bot_config)
          for bot_name, bot_config in config.bots.items()
      },
      population_size=config.num_bots,
  )

  return scenario_lib.Scenario(
      substrate=substrate,<commit_msg>Allow a custom transform to be applied to a scenario's underlying substrate.<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.
"""Scenario factory."""

from typing import Collection, Mapping

from ml_collections import config_dict

from meltingpot.python import bot as bot_factory
from meltingpot.python import substrate as substrate_factory
from meltingpot.python.configs import scenarios as scenario_config
from meltingpot.python.utils.scenarios import population
from meltingpot.python.utils.scenarios import scenario as scenario_lib
from meltingpot.python.utils.scenarios import substrate_transforms

AVAILABLE_SCENARIOS = frozenset(scenario_config.SCENARIO_CONFIGS)

SCENARIOS_BY_SUBSTRATE: Mapping[
    str, Collection[str]] = scenario_config.scenarios_by_substrate(
        scenario_config.SCENARIO_CONFIGS)

PERMITTED_OBSERVATIONS = frozenset({
    'INVENTORY',
    'READY_TO_SHOOT',
    'RGB',
})

# TODO(b/227143834): Remove aliases once internal deps have been removed.
Scenario = scenario_lib.Scenario
PopulationObservables = population.PopulationObservables
ScenarioObservables = scenario_lib.ScenarioObservables
Population = population.Population


def get_config(scenario_name: str) -> config_dict.ConfigDict:
  """Returns a config for the specified scenario.

  Args:
    scenario_name: Name of the scenario. Must be in AVAILABLE_SCENARIOS.
  """
  if scenario_name not in AVAILABLE_SCENARIOS:
    raise ValueError(f'Unknown scenario {scenario_name!r}')
  scenario = scenario_config.SCENARIO_CONFIGS[scenario_name]
  substrate = substrate_factory.get_config(scenario.substrate)
  bots = {name: bot_factory.get_config(name) for name in scenario.bots}
  config = config_dict.create(
      substrate=substrate,
      bots=bots,
      is_focal=scenario.is_focal,
      num_players=sum(scenario.is_focal),
      num_bots=len(scenario.is_focal) - sum(scenario.is_focal),
      substrate_transform=lambda x: x,
      permitted_observations=set(PERMITTED_OBSERVATIONS),
  )
  return config.lock()


def build(config: config_dict.ConfigDict) -> scenario_lib.Scenario:
  """Builds a scenario for the given config.

  Args:
    config: config resulting from `get_config`.

  Returns:
    The test scenario.
  """
  substrate = substrate_factory.build(config.substrate)
  substrate: substrate_factory.Substrate = config.substrate_transform(substrate)
  # Add observations needed by some bots. These are removed for focal players.
  permitted_observations = (
      set(substrate.observation_spec()[0]) & config.permitted_observations)
  substrate = substrate_transforms.with_tf1_bot_required_observations(substrate)

  background_population = population.Population(
      policies={
          bot_name: bot_factory.build(bot_config)
          for bot_name, bot_config in config.bots.items()
      },
      population_size=config.num_bots,
  )

  return scenario_lib.Scenario(
      substrate=substrate,
{'fraction_non_alphanumeric': 0.07573716509504565, 'fraction_numerical': 0.005687771291722796, 'mean_word_length': 4.8256320836966, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '1658517', 'n_tokens_mistral': 2053, 'n_tokens_neox': 1981, 'n_words': 541}