Merge pull request #13 from zbeucler2018/develop

Author: zbeucler2018

.gitmodules

@@ -20,7 +20,7 @@
       "address": 33583330,
       "type": "|i1"
     },
-    "hit wall?": {
+    "hit_wall": {
       "address": 33583328,
       "type": "><n4"
     }

HotWheelsEnv.py

@@ -1,34 +1,41 @@
 
 -- By: Zack Beucler
+-- single lap: (data.progress == 320)
+-- three laps: (data.progress >= 950)
 
-function isDone()
-	-- local gameover_val = 33583680
-	-- if data.score > gameover_val then
-	-- 	return true
-	-- else
-	-- 	return false
-	-- end
-
-	local single_lap = 320
-	local three_laps = 950
-	local LAP_LIMIT = single_lap
-	if data.progress >= LAP_LIMIT then
+
+function isGameOver()
+	-- if data.progress >= 949 then
+	if data.lap >= 4 then
+		return true
+	else
+		return false
+	end
+end
+
+
+function isHittingWall()
+	if data.hit_wall > 100 then
 		return true
 	else
 		return false
 	end
 end
 
 
+function isDone()
+	return isGameOver() or isHittingWall()
+end
+
+
 previous_progress = 0
 function calculateReward()
 	local current_progress = data.progress
+	local delta = 0
 	if current_progress > previous_progress then
-		local delta = current_progress - previous_progress
+		delta = current_progress - previous_progress
 		previous_progress = current_progress
-		return delta
-	else
-		return 0
 	end
+	return delta
 end
 

HotWheelsStuntTrackChallenge-gba/data.json

@@ -0,0 +1,28 @@
+import pprint
+
+import gymnasium as gym
+from stable_baselines3.common.callbacks import BaseCallback, CallbackList, EvalCallback
+from wandb.integration.sb3 import WandbCallback
+
+import wandb
+
+# log to wandb
+
+# run agent and record video?
+
+
+class HotWheelsCallback(WandbCallback):
+    def __init__(verbose, _model_save_path, _model_save_freq):
+        super().__init__(
+            verbose=1,
+            model_save_path=_model_save_path,
+            model_save_freq=_model_save_freq,
+        )
+
+    def _on_step(self) -> bool:
+        # print(self.locals['infos'])
+        for cpu in self.locals["infos"]:
+            print(cpu)
+            wandb.log(cpu)
+        # wandb.log(self.locals['infos'][0])
+        return super()._on_step()

HotWheelsStuntTrackChallenge-gba/script.lua

@@ -0,0 +1,144 @@
+from enum import Enum
+from typing import Any, Tuple, Union
+
+import retro
+from gymnasium.core import Env
+from gymnasium.wrappers import GrayScaleObservation, ResizeObservation
+
+
+from gym_wrappers import (
+    LogInfoValues,
+    NorrmalizeBoost,
+    PunishHittingWalls,
+    EncourageTricks,
+    FixSpeed,
+    TerminateOnCrash,
+    HotWheelsDiscretizer,
+    CropObservation
+)
+
+
+import os
+from typing import Any, Callable, Dict, Optional, Type, Union
+
+from stable_baselines3.common.monitor import Monitor
+from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecEnv
+
+
+class GameStates(Enum):
+    """
+    Possible game states
+    """
+
+    SINGLE = "dino_single.state"
+    SINGLE_POINTS = "dino_single_points.state"
+    MULTIPLAYER = "dino_multiplayer.state"
+
+
+def make_hotwheels_vec_env(
+    env_id: Union[str, Callable[..., Env]],
+    game_state: str,
+    n_envs: int = 1,
+    seed: Optional[int] = None,
+    start_index: int = 0,
+    monitor_dir: Optional[str] = None,
+    wrapper_class: Optional[Callable[[Env], Env]] = None,
+    env_kwargs: Optional[Dict[str, Any]] = None,
+    vec_env_cls: Optional[Type[Union[DummyVecEnv, SubprocVecEnv]]] = None,
+    extra_wrappers: Optional[list[Env]] = None,
+    vec_env_kwargs: Optional[Dict[str, Any]] = None,
+    monitor_kwargs: Optional[Dict[str, Any]] = None,
+    wrapper_kwargs: Optional[Dict[str, Any]] = None,
+) -> VecEnv:
+    """
+    Create a wrapped, monitored ``VecEnv``.
+    By default it uses a ``DummyVecEnv`` which is usually faster
+    than a ``SubprocVecEnv``.
+    Modified for HotWheels
+
+    :param env_id: either the env ID, the env class or a callable returning an env
+    :param n_envs: the number of environments you wish to have in parallel
+    :param seed: the initial seed for the random number generator
+    :param start_index: start rank index
+    :param monitor_dir: Path to a folder where the monitor files will be saved.
+        If None, no file will be written, however, the env will still be wrapped
+        in a Monitor wrapper to provide additional information about training.
+    :param wrapper_class: Additional wrapper to use on the environment.
+        This can also be a function with single argument that wraps the environment in many things.
+        Note: the wrapper specified by this parameter will be applied after the ``Monitor`` wrapper.
+        if some cases (e.g. with TimeLimit wrapper) this can lead to undesired behavior.
+        See here for more details: https://github.com/DLR-RM/stable-baselines3/issues/894
+    :param env_kwargs: Optional keyword argument to pass to the env constructor
+    :param vec_env_cls: A custom ``VecEnv`` class constructor. Default: None.
+    :param extra_wrappers: Optional list of wrappers to wrap the env. Applied after using the ``Monitor`` wrapper.
+    :param vec_env_kwargs: Keyword arguments to pass to the ``VecEnv`` class constructor.
+    :param monitor_kwargs: Keyword arguments to pass to the ``Monitor`` class constructor.
+    :param wrapper_kwargs: Keyword arguments to pass to the ``Wrapper`` class constructor.
+    :return: The wrapped environment
+    """
+    env_kwargs = env_kwargs or {}
+    vec_env_kwargs = vec_env_kwargs or {}
+    monitor_kwargs = monitor_kwargs or {}
+    wrapper_kwargs = wrapper_kwargs or {}
+    assert vec_env_kwargs is not None  # for mypy
+
+    def make_env(rank: int) -> Callable[[], Env]:
+        def _init() -> Env:
+            # For type checker:
+            assert monitor_kwargs is not None
+            assert wrapper_kwargs is not None
+            assert env_kwargs is not None
+
+            if isinstance(env_id, str):
+                # if the render mode was not specified, we set it to `rgb_array` as default.
+                kwargs = {"render_mode": "rgb_array"}
+                kwargs.update(env_kwargs)
+                env = retro.make(env_id, state=game_state, **kwargs)  # type: ignore[arg-type]
+                env = TerminateOnCrash(env)
+                env = FixSpeed(env)
+                env = EncourageTricks(env)
+                env = HotWheelsDiscretizer(env)
+                env = CropObservation(env)
+                env = ResizeObservation(env, (84, 84))
+                # env = LogInfoValues(env)
+            else:
+                env = env_id(**env_kwargs)
+                # Patch to support gym 0.21/0.26 and gymnasium
+                # env = _patch_env(env)
+
+            if seed is not None:
+                # Note: here we only seed the action space
+                # We will seed the env at the next reset
+                env.action_space.seed(seed + rank)
+            # Wrap the env in a Monitor wrapper
+            # to have additional training information
+            monitor_path = (
+                os.path.join(monitor_dir, str(rank))
+                if monitor_dir is not None
+                else None
+            )
+            # Create the monitor folder if needed
+            if monitor_path is not None and monitor_dir is not None:
+                os.makedirs(monitor_dir, exist_ok=True)
+            env = Monitor(env, filename=monitor_path, **monitor_kwargs)
+            # Optionally, wrap the environment with the provided wrapper
+            if wrapper_class is not None:
+                env = wrapper_class(env, **wrapper_kwargs)
+            if extra_wrappers is not None:
+                for _wrapper in extra_wrappers:
+                    env = _wrapper(extra_wrappers)
+            return env
+
+        return _init
+
+    # No custom VecEnv is passed
+    if vec_env_cls is None:
+        # Default: use a DummyVecEnv
+        vec_env_cls = DummyVecEnv
+
+    vec_env = vec_env_cls(
+        [make_env(i + start_index) for i in range(n_envs)], **vec_env_kwargs
+    )
+    # Prepare the seeds for the first reset
+    vec_env.seed(seed)
+    return vec_env