adding the wrapper and others

junkilee · Apr 1, 2019 · 03c09c2 · 03c09c2
1 parent 3e499a4
commit 03c09c2
Show file tree

Hide file tree

Showing 17 changed files with 580 additions and 212 deletions.
diff --git a/.idea/starcraft-ai.iml b/.idea/starcraft-ai.iml
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
@@ -32,8 +32,11 @@ Install pysc2
    - http://github.com/junkilee/pysc2
    - `pip install -e .`
 
-## Install docker
+Install docker
+-------------
  - Installation instructions for Ubuntu
    - https://docs.docker.com/install/linux/docker-ce/ubuntu/
  - Post installation setup
    - https://docs.docker.com/install/linux/linux-postinstall/
+ - Nvidia-docker setup
+   - https://github.com/NVIDIA/nvidia-docker
diff --git a/sc2ai/envs/actions.py b/sc2ai/envs/actions.py
diff --git a/sc2ai/envs/minigame.py b/sc2ai/envs/minigame.py
diff --git a/sc2ai/envs/minigames/__init__.py b/sc2ai/envs/minigames/__init__.py
@@ -0,0 +1,5 @@
+from sc2ai.envs.minigames.collect_minerals_and_gas import CollectMineralAndGasEnv
+from sc2ai.envs.minigames.defeat_zerglings import DefeatZerglingsEnv
+from sc2ai.envs.minigames.move_to_beacon import MoveToBeaconEnv
+
+__all__ = ["CollectMineralAndGasEnv", "DefeatZerglingsEnv", "MoveToBeaconEnv"]
diff --git a/sc2ai/envs/minigames/collect_minerals_and_gas.py b/sc2ai/envs/minigames/collect_minerals_and_gas.py
@@ -2,9 +2,9 @@
 from gym import spaces
 from pysc2.lib import actions, features
 
-from ..minigame import MiniGameEnv
+from ..sc2env import MiniGameEnv
 
-class DefeatZerglingsEnv(MiniGameEnv):
+class CollectMineralAndGasEnv(MiniGameEnv):
     """
 
     """

diff --git a/sc2ai/envs/minigames/defeat_zerglings.py b/sc2ai/envs/minigames/defeat_zerglings.py
@@ -2,9 +2,12 @@
 from gym import spaces
 from pysc2.lib import actions, features
 
-from ..minigame import MiniGameEnv
+from ..sc2env import MiniGameEnv
 
 class DefeatZerglingsEnv(MiniGameEnv):
+    """
+
+    """
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 

diff --git a/sc2ai/envs/minigames/move_to_beacon.py b/sc2ai/envs/minigames/move_to_beacon.py
@@ -2,12 +2,14 @@
 from gym import spaces
 from pysc2.lib import actions, features
 
-from ..minigame import MiniGameEnv
+from ..sc2env import MiniGameEnv
 
-class MoveToBeaconDiscreteEnv(MiniGameEnv):
+class MoveToBeaconEnv(MiniGameEnv):
+    """
+
+    """
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-
 
     def _process_reward(self, reward, raw_obs):
         raise NotImplementedError

diff --git a/sc2ai/envs/sc2env.py b/sc2ai/envs/sc2env.py
@@ -11,11 +11,35 @@
 
 from .game_info import default_env_options
 from .game_info import ActionIDs
+import sc2ai.envs.minigames as minigames
 
 import numpy as np
 
 from baselines import logger
 
+MAP_ENV_MAPPINGS = {
+    "DefeatZerglings": minigames.DefeatZerglingsEnv,
+    "CollectMineralAndGas": minigames.DefeatCollectMinardAndGasEnv,
+    "MoveToBeacon": minigames.MoveToBeaconEnv
+}
+
+def make_sc2env(**kwargs):
+    """
+
+    Args:
+        **kwargs:
+
+    Returns:
+
+    """
+    cls = None
+    if kwargs["map"] is not in MAP_ENV_MAPPINGS:
+        raise Exception("The map is unknown and not registered.")
+    else
+        cls = MAP_ENV_MAPPINGS[kwargs["map"]]
+
+    return gym.make(cls(**kwargs))
+
 class SC2Env(gym.Env):
     """A gym wrapper for PySC2's Starcraft II environment.
 
@@ -136,7 +160,7 @@ def step(self, action):
             logging.info("Keyboard Interruption.")
             return None, 0, True, {}
         except Exception:
-            logger.exception("An unexpected exception occured.")
+            logging.exception("An unexpected exception occured.")
             return None, 0, True, {}
         self._available_actions = obs.observation['available_actions']
         reward = obs.reward
@@ -155,4 +179,49 @@ def action_space(self):
 
     @property
     def observation_space(self):
-        raise self.observation_space
+        raise self.observation_space
+
+
+class MiniGameEnv(SC2Env):
+    """Providing a wrapper for minigames. Mainly supports preprocessing of both reward and observation.
+
+    Args:
+        map_name:
+        **kwargs:
+    """
+    def __init__(self, map_name, **kwargs):
+        assert isinstance(map_name, str)
+        super().__init__(map_name, **kwargs)
+
+    def _reset(self):
+        obs = super()._reset()
+        return _process_observation(self, obs)
+
+    def _step(self, action):
+        raw_obs, reward, done, info = super()._step(self, action)
+        reward = self._process_reward(reward, raw_obs)
+        obs = self._process_observation(raw_obs)
+        return obs, reward, done, info
+
+    def _process_reward(self, reward, raw_obs):
+        raise NotImplementedError
+
+    def _process_observation(self, raw_obs):
+        raise NotImplementedError
+
+class MultiStepMiniGameEnv(SC2Env):
+    """
+    Instead of taking one action per step, it handles a list of consecutive actions.
+    """
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def _step(self, actions):
+        total_reward = 0
+        for action in actions:
+            raw_obs, reward, done, info = super().__step(self, action)
+            total_reward += self._process_reward(reward, raw_obs)
+            if done:
+                break
+        obs = self._process_observation(raw_obs)
+        return obs, total_reward, done, info
diff --git a/sc2ai/policies/ray_fully_conv_lstm.py → sc2ai/models/rllib_models.py b/sc2ai/policies/ray_fully_conv_lstm.py → sc2ai/models/rllib_models.py
diff --git a/sc2ai/policies/bl_fully_conv_lstm.py b/sc2ai/policies/bl_fully_conv_lstm.py