mirror of
https://github.com/gryf/coach.git
synced 2025-12-18 03:30:19 +01:00
pre-release 0.10.0
This commit is contained in:
Gal Novik
245
rl_coach/environments/starcraft2_environment.py
Normal file
245
rl_coach/environments/starcraft2_environment.py
Normal file
@@ -0,0 +1,245 @@
|
||||
#
|
||||
# Copyright (c) 2017 Intel Corporation
|
||||
#
|
||||
# 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
|
||||
#
|
||||
# http://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.
|
||||
#
|
||||
|
||||
|
||||
from enum import Enum
|
||||
from typing import Union, List
|
||||
|
||||
import numpy as np
|
||||
from rl_coach.filters.observation.observation_move_axis_filter import ObservationMoveAxisFilter
|
||||
|
||||
try:
|
||||
from pysc2 import maps
|
||||
from pysc2.env import sc2_env
|
||||
from pysc2.env import available_actions_printer
|
||||
from pysc2.lib import actions
|
||||
from pysc2.lib import features
|
||||
from pysc2.env import environment
|
||||
from absl import app
|
||||
from absl import flags
|
||||
except ImportError:
|
||||
from rl_coach.logger import failed_imports
|
||||
failed_imports.append("PySc2")
|
||||
|
||||
from rl_coach.environments.environment import Environment, EnvironmentParameters, LevelSelection
|
||||
from rl_coach.base_parameters import VisualizationParameters
|
||||
from rl_coach.spaces import BoxActionSpace, VectorObservationSpace, PlanarMapsObservationSpace, StateSpace, CompoundActionSpace, \
|
||||
DiscreteActionSpace
|
||||
from rl_coach.filters.filter import InputFilter, OutputFilter
|
||||
from rl_coach.filters.observation.observation_rescale_to_size_filter import ObservationRescaleToSizeFilter
|
||||
from rl_coach.filters.action.linear_box_to_box_map import LinearBoxToBoxMap
|
||||
from rl_coach.filters.observation.observation_to_uint8_filter import ObservationToUInt8Filter
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
FLAGS(['coach.py'])
|
||||
|
||||
SCREEN_SIZE = 84 # will also impact the action space size
|
||||
|
||||
# Starcraft Constants
|
||||
_NOOP = actions.FUNCTIONS.no_op.id
|
||||
_MOVE_SCREEN = actions.FUNCTIONS.Move_screen.id
|
||||
_SELECT_ARMY = actions.FUNCTIONS.select_army.id
|
||||
_PLAYER_RELATIVE = features.SCREEN_FEATURES.player_relative.index
|
||||
_NOT_QUEUED = [0]
|
||||
_SELECT_ALL = [0]
|
||||
|
||||
|
||||
class StarcraftObservationType(Enum):
|
||||
Features = 0
|
||||
RGB = 1
|
||||
|
||||
|
||||
StarcraftInputFilter = InputFilter(is_a_reference_filter=True)
|
||||
StarcraftInputFilter.add_observation_filter('screen', 'move_axis', ObservationMoveAxisFilter(0, -1))
|
||||
StarcraftInputFilter.add_observation_filter('screen', 'rescaling',
|
||||
ObservationRescaleToSizeFilter(
|
||||
PlanarMapsObservationSpace(np.array([84, 84, 1]),
|
||||
low=0, high=255, channels_axis=-1)))
|
||||
StarcraftInputFilter.add_observation_filter('screen', 'to_uint8', ObservationToUInt8Filter(0, 255))
|
||||
|
||||
StarcraftInputFilter.add_observation_filter('minimap', 'move_axis', ObservationMoveAxisFilter(0, -1))
|
||||
StarcraftInputFilter.add_observation_filter('minimap', 'rescaling',
|
||||
ObservationRescaleToSizeFilter(
|
||||
PlanarMapsObservationSpace(np.array([64, 64, 1]),
|
||||
low=0, high=255, channels_axis=-1)))
|
||||
StarcraftInputFilter.add_observation_filter('minimap', 'to_uint8', ObservationToUInt8Filter(0, 255))
|
||||
|
||||
|
||||
StarcraftNormalizingOutputFilter = OutputFilter(is_a_reference_filter=True)
|
||||
StarcraftNormalizingOutputFilter.add_action_filter(
|
||||
'normalization', LinearBoxToBoxMap(input_space_low=-SCREEN_SIZE / 2, input_space_high=SCREEN_SIZE / 2 - 1))
|
||||
|
||||
|
||||
class StarCraft2EnvironmentParameters(EnvironmentParameters):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.screen_size = 84
|
||||
self.minimap_size = 64
|
||||
self.feature_minimap_maps_to_use = range(7)
|
||||
self.feature_screen_maps_to_use = range(17)
|
||||
self.observation_type = StarcraftObservationType.Features
|
||||
self.disable_fog = False
|
||||
self.auto_select_all_army = True
|
||||
self.default_input_filter = StarcraftInputFilter
|
||||
self.default_output_filter = StarcraftNormalizingOutputFilter
|
||||
self.use_full_action_space = False
|
||||
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return 'rl_coach.environments.starcraft2_environment:StarCraft2Environment'
|
||||
|
||||
|
||||
# Environment
|
||||
class StarCraft2Environment(Environment):
|
||||
def __init__(self, level: LevelSelection, frame_skip: int, visualization_parameters: VisualizationParameters,
|
||||
seed: Union[None, int]=None, human_control: bool=False,
|
||||
custom_reward_threshold: Union[int, float]=None,
|
||||
screen_size: int=84, minimap_size: int=64,
|
||||
feature_minimap_maps_to_use: List=range(7), feature_screen_maps_to_use: List=range(17),
|
||||
observation_type: StarcraftObservationType=StarcraftObservationType.Features,
|
||||
disable_fog: bool=False, auto_select_all_army: bool=True,
|
||||
use_full_action_space: bool=False, **kwargs):
|
||||
super().__init__(level, seed, frame_skip, human_control, custom_reward_threshold, visualization_parameters)
|
||||
|
||||
self.screen_size = screen_size
|
||||
self.minimap_size = minimap_size
|
||||
self.feature_minimap_maps_to_use = feature_minimap_maps_to_use
|
||||
self.feature_screen_maps_to_use = feature_screen_maps_to_use
|
||||
self.observation_type = observation_type
|
||||
self.features_screen_size = None
|
||||
self.feature_minimap_size = None
|
||||
self.rgb_screen_size = None
|
||||
self.rgb_minimap_size = None
|
||||
if self.observation_type == StarcraftObservationType.Features:
|
||||
self.features_screen_size = screen_size
|
||||
self.feature_minimap_size = minimap_size
|
||||
elif self.observation_type == StarcraftObservationType.RGB:
|
||||
self.rgb_screen_size = screen_size
|
||||
self.rgb_minimap_size = minimap_size
|
||||
self.disable_fog = disable_fog
|
||||
self.auto_select_all_army = auto_select_all_army
|
||||
self.use_full_action_space = use_full_action_space
|
||||
|
||||
# step_mul is the equivalent to frame skipping. Not sure if it repeats actions in between or not though.
|
||||
self.env = sc2_env.SC2Env(map_name=self.env_id, step_mul=frame_skip,
|
||||
visualize=self.is_rendered,
|
||||
agent_interface_format=sc2_env.AgentInterfaceFormat(
|
||||
feature_dimensions=sc2_env.Dimensions(
|
||||
screen=self.features_screen_size,
|
||||
minimap=self.feature_minimap_size
|
||||
)
|
||||
# rgb_dimensions=sc2_env.Dimensions(
|
||||
# screen=self.rgb_screen_size,
|
||||
# minimap=self.rgb_screen_size
|
||||
# )
|
||||
),
|
||||
# feature_screen_size=self.features_screen_size,
|
||||
# feature_minimap_size=self.feature_minimap_size,
|
||||
# rgb_screen_size=self.rgb_screen_size,
|
||||
# rgb_minimap_size=self.rgb_screen_size,
|
||||
disable_fog=disable_fog,
|
||||
random_seed=self.seed
|
||||
)
|
||||
|
||||
# print all the available actions
|
||||
# self.env = available_actions_printer.AvailableActionsPrinter(self.env)
|
||||
|
||||
self.reset_internal_state(True)
|
||||
|
||||
"""
|
||||
feature_screen: [height_map, visibility_map, creep, power, player_id, player_relative, unit_type, selected,
|
||||
unit_hit_points, unit_hit_points_ratio, unit_energy, unit_energy_ratio, unit_shields,
|
||||
unit_shields_ratio, unit_density, unit_density_aa, effects]
|
||||
feature_minimap: [height_map, visibility_map, creep, camera, player_id, player_relative, selecte
|
||||
d]
|
||||
player: [player_id, minerals, vespene, food_cap, food_army, food_workers, idle_worker_dount,
|
||||
army_count, warp_gate_count, larva_count]
|
||||
"""
|
||||
self.screen_shape = np.array(self.env.observation_spec()[0]['feature_screen'])
|
||||
self.screen_shape[0] = len(self.feature_screen_maps_to_use)
|
||||
self.minimap_shape = np.array(self.env.observation_spec()[0]['feature_minimap'])
|
||||
self.minimap_shape[0] = len(self.feature_minimap_maps_to_use)
|
||||
self.state_space = StateSpace({
|
||||
"screen": PlanarMapsObservationSpace(shape=self.screen_shape, low=0, high=255, channels_axis=0),
|
||||
"minimap": PlanarMapsObservationSpace(shape=self.minimap_shape, low=0, high=255, channels_axis=0),
|
||||
"measurements": VectorObservationSpace(self.env.observation_spec()[0]["player"][0])
|
||||
})
|
||||
if self.use_full_action_space:
|
||||
action_identifiers = list(self.env.action_spec()[0].functions)
|
||||
num_action_identifiers = len(action_identifiers)
|
||||
action_arguments = [(arg.name, arg.sizes) for arg in self.env.action_spec()[0].types]
|
||||
sub_action_spaces = [DiscreteActionSpace(num_action_identifiers)]
|
||||
for argument in action_arguments:
|
||||
for dimension in argument[1]:
|
||||
sub_action_spaces.append(DiscreteActionSpace(dimension))
|
||||
self.action_space = CompoundActionSpace(sub_action_spaces)
|
||||
else:
|
||||
self.action_space = BoxActionSpace(2, 0, self.screen_size - 1, ["X-Axis, Y-Axis"],
|
||||
default_action=np.array([self.screen_size/2, self.screen_size/2]))
|
||||
|
||||
def _update_state(self):
|
||||
timestep = 0
|
||||
self.screen = self.last_result[timestep].observation.feature_screen
|
||||
# extract only the requested segmentation maps from the observation
|
||||
self.screen = np.take(self.screen, self.feature_screen_maps_to_use, axis=0)
|
||||
self.minimap = self.last_result[timestep].observation.feature_minimap
|
||||
self.measurements = self.last_result[timestep].observation.player
|
||||
self.reward = self.last_result[timestep].reward
|
||||
self.done = self.last_result[timestep].step_type == environment.StepType.LAST
|
||||
self.state = {
|
||||
'screen': self.screen,
|
||||
'minimap': self.minimap,
|
||||
'measurements': self.measurements
|
||||
}
|
||||
|
||||
def _take_action(self, action):
|
||||
if self.use_full_action_space:
|
||||
action_identifier = action[0]
|
||||
action_arguments = action[1:]
|
||||
action = actions.FunctionCall(action_identifier, action_arguments)
|
||||
else:
|
||||
coord = np.array(action[0:2])
|
||||
noop = False
|
||||
coord = coord.round()
|
||||
coord = np.clip(coord, 0, SCREEN_SIZE - 1)
|
||||
self.last_action_idx = coord
|
||||
|
||||
if noop:
|
||||
action = actions.FunctionCall(_NOOP, [])
|
||||
else:
|
||||
action = actions.FunctionCall(_MOVE_SCREEN, [_NOT_QUEUED, coord])
|
||||
|
||||
self.last_result = self.env.step(actions=[action])
|
||||
|
||||
def _restart_environment_episode(self, force_environment_reset=False):
|
||||
# reset the environment
|
||||
self.last_result = self.env.reset()
|
||||
|
||||
# select all the units on the screen
|
||||
if self.auto_select_all_army:
|
||||
self.env.step(actions=[actions.FunctionCall(_SELECT_ARMY, [_SELECT_ALL])])
|
||||
|
||||
def get_rendered_image(self):
|
||||
screen = np.squeeze(np.tile(np.expand_dims(self.screen, -1), (1, 1, 3)))
|
||||
screen = screen / np.max(screen) * 255
|
||||
return screen.astype('uint8')
|
||||
|
||||
def dump_video_of_last_episode(self):
|
||||
from rl_coach.logger import experiment_path
|
||||
self.env._run_config.replay_dir = experiment_path
|
||||
self.env.save_replay('replays')
|
||||
super().dump_video_of_last_episode()
|
||||
Reference in New Issue
Block a user