genesis-forge 0.0.1__py3-none-any.whl

genesis_forge/__init__.py

@@ -0,0 +1,12 @@
+from .rl.skrl import create_skrl_env
+from .wrappers import VideoWrapper
+from .genesis_env import GenesisEnv, EnvMode
+from .managed_env import ManagedEnvironment
+
+__all__ = [
+    "GenesisEnv",
+    "ManagedEnvironment",
+    "EnvMode",
+    "VideoWrapper",
+    "create_skrl_env",
+]

genesis_forge/gamepads/__init__.py

@@ -0,0 +1,5 @@
+from .gamepad import Gamepad
+from .logitech import LOGITECH_F710_CONFIG, LOGITECH_F310_CONFIG
+from .config import GamepadState
+
+__all__ = ["Gamepad", "GamepadState", "LOGITECH_F710_CONFIG", "LOGITECH_F310_CONFIG"]

genesis_forge/gamepads/config.py

@@ -0,0 +1,104 @@
+from typing import TypedDict
+from enum import Enum
+
+
+class Button(Enum):
+    A = 1
+    B = 2
+    X = 3
+    Y = 4
+    LB = 5
+    RB = 6
+    LT = 7
+    RT = 8
+    BACK = 9
+    START = 10
+    MODE = 11
+    UP = 12
+    DOWN = 13
+    LEFT = 14
+    RIGHT = 15
+    LEFT_JOYSTICK = 16
+    RIGHT_JOYSTICK = 17
+
+
+class DPad(Enum):
+    NONE = 0
+    UP = 1
+    DOWN = 2
+    LEFT = 3
+    RIGHT = 4
+
+
+class GamepadMapping(TypedDict):
+    """
+    Defines the how to extract a value from a gamepad data array.
+    """
+
+    data: int
+    """The index of the data in the gamepad data array."""
+
+    button: Button
+    """If this is a button, what is it's name"""
+
+    axis: int
+    """If this is a joystick, what is it's axis number"""
+
+    dpad: DPad
+    """D-pad direction."""
+
+    bitmask: int
+    """The bitmask to extract the value from the data."""
+
+    value: int
+    """Match this value to the value at the data index."""
+
+
+class GamepadState:
+    """Data from a gamepad."""
+
+    axis_values: list[float] = []
+    """The value (-1 to 1) for each axes."""
+
+    buttons: list[str] = []
+    """A list of the buttons that are pressed."""
+
+    dpad: DPad
+    """D-pad direction."""
+
+    def axis(self, index: int):
+        """
+        Get the axis value at an index.
+        This is the preferred way to get the axis value, because the axis array will not be filled until the gamepad
+        receives input.
+
+        Args:
+            index: The index of the axis to get the value of.
+
+        Returns:
+            The value of the axis at the index.
+        """
+        if index >= len(self.axis_values):
+            return 0.0
+        return self.axis_values[index]
+
+    def __repr__(self):
+        return f"GamepadState(axis={self.axis_values}, buttons={self.buttons}, dpad={self.dpad})"
+
+
+class GamepadConfig(TypedDict):
+    """
+    Defines a gamepad, how to connect to it, and how to map the buttons and axes.
+    """
+
+    name: str
+    """The name of the gamepad."""
+
+    vendor_id: int
+    """The vendor id of the gamepad."""
+
+    product_id: int
+    """The product id of the gamepad."""
+
+    mapping: list[GamepadMapping]
+    """The mapping of the gamepad."""

genesis_forge/gamepads/debug.py

@@ -0,0 +1,34 @@
+import time
+import argparse
+
+from .gamepad import Gamepad, GamepadState
+
+
+class DebugGamepad(Gamepad):
+    """
+    This is just used to output the HID data to the console so you can implement a new gamepad class.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def parse_data(self, data) -> GamepadState:
+        print(data)
+        # print([bin(d) for d in data])
+        return GamepadState()
+
+
+if __name__ == "__main__":
+    """Run from the CLI with python -m genesis_forge.managers.command.gamepads.debug --help"""
+    parser = argparse.ArgumentParser(add_help=True)
+    parser.add_argument("-v", "--vender_id", type=str)
+    parser.add_argument("-p", "--product_id", type=str)
+    args = parser.parse_args()
+
+    vendor_id = int(args.vender_id, 16)
+    product_id = int(args.product_id, 16)
+
+    gamepad = DebugGamepad(vendor_id=vendor_id, product_id=product_id)
+    while True:
+        # print(gamepad.get_command())
+        time.sleep(0.1)

genesis_forge/gamepads/gamepad.py

@@ -0,0 +1,191 @@
+"""Logitech F310/F710 Gamepad class that uses HID under the hood.
+
+Adapted from: https://github.com/google-deepmind/mujoco_playground/blob/a873d53765a4c83572cf44fa74768ab62ceb7be1/mujoco_playground/experimental/sim2sim/gamepad_reader.py.
+"""
+
+import time
+import argparse
+import hid
+import threading
+
+from .config import GamepadConfig, GamepadState
+from .logitech import LOGITECH_F710_CONFIG, LOGITECH_F310_CONFIG
+
+GAMEPAD_CONFIGS = [
+    LOGITECH_F710_CONFIG,
+    LOGITECH_F310_CONFIG,
+]
+
+
+class Gamepad:
+    """
+    General gamepad controller.
+    If a config is not provided, it will automatically attempt to connect to one of the configured gamepads.
+
+    Example:
+        >>> gamepad = Gamepad(config=LOGITECH_F710_CONFIG)
+        >>> gamepad.state
+        GamepadState(axis=[0.0, 0.0, 0.0, 0.0], buttons=[A], dpad=UP)
+        >>> gamepad.state.axis
+        [0.0, 0.0, 0.0, 0.0]
+        >>> gamepad.state.buttons
+        ["A"]
+        >>> gamepad.state.dpad
+        "UP"
+        >>> gamepad.state.buttons = [Button.A]
+    """
+
+    def __init__(
+        self,
+        config: GamepadConfig = None,
+        vendor_id=None,
+        product_id=None,
+        debug=False,
+    ):
+        self._config = config
+        self._vendor_id = vendor_id
+        self._product_id = product_id
+
+        if vendor_id is None and config is not None:
+            self._vendor_id = config["vendor_id"]
+        if product_id is None and config is not None:
+            self._product_id = config["product_id"]
+
+        self._state = GamepadState()
+        self._debug = debug
+
+        self.is_running = True
+        self._device = None
+
+        self.connect()
+        self.read_thread = threading.Thread(target=self.read_loop, daemon=True)
+        self.read_thread.start()
+
+    @property
+    def state(self) -> GamepadState:
+        """
+        The current state of the gamepad.
+        """
+        return self._state
+
+    def parse_data(self, data: list[int]) -> GamepadState:
+        """Parse gamepad data into a GamepadState object."""
+        axis = []
+        buttons = []
+        dpad = None
+
+        # No gamepad config, so we cann't parse the data
+        if self._config is None:
+            return
+
+        for cfg in self._config["mapping"]:
+            if "data" not in cfg:
+                print(f"Warning: {cfg} has no data value")
+                continue
+            if cfg["data"] >= len(data):
+                print(f"Error: {cfg} data is out of range")
+                continue
+            value = data[cfg["data"]]
+            value_truthy = False
+
+            # Apply the bitmask to the value
+            if "bitmask" in cfg:
+                value = value & cfg["bitmask"]
+                if value != 0:
+                    value_truthy = True
+            elif "button" in cfg or "dpad" in cfg:
+                print(f"Warning: {cfg} has no bitmask value")
+                continue
+
+            # Check if value is matches
+            if "value" in cfg:
+                value_truthy = value == cfg["value"]
+
+            if "button" in cfg and value_truthy:
+                buttons.append(cfg["button"].name)
+            elif "dpad" in cfg and value_truthy:
+                dpad = cfg["dpad"].name
+            elif "axis" in cfg:
+                value = -(value - 128) / 128.0
+                axis.insert(cfg["axis"], value)
+
+        self._state.axis_values = axis
+        self._state.buttons = buttons
+        self._state.dpad = dpad
+        return self._state
+
+    def auto_connect(self):
+        """Loop through the available gamepad configs until one connects."""
+        for config in GAMEPAD_CONFIGS:
+            self._vendor_id = config["vendor_id"]
+            self._product_id = config["product_id"]
+            self._config = config
+            try:
+                if self.connect():
+                    return
+            except:
+                pass
+        raise IOError(f"Could not find a gamepad to connect to")
+
+    def connect(self, vendor_id=None, product_id=None):
+        """
+        Attempt to connect to a gamepad.
+        """
+        if vendor_id is None:
+            vendor_id = self._vendor_id
+        if product_id is None:
+            product_id = self._product_id
+
+        # If the vendor/product IDs aren't set, loop through the available gamepad configs
+        if product_id is None and vendor_id is None:
+            self.auto_connect()
+            return
+
+        try:
+            self._device = hid.device()
+            self._device.open(vendor_id, product_id)
+            self._device.set_nonblocking(True)
+            print(
+                f"Connected to gamepad {self._device.get_manufacturer_string()} {self._device.get_product_string()}"
+            )
+            return True
+        except IOError as e:
+            raise IOError(
+                f"Error connecting to gamepad 0x{vendor_id:04x}:0x{product_id:04x}: {e}"
+            )
+
+    def read_loop(self):
+        """
+        Wait for gamepad input, and then update the gamepad state.
+        """
+        while self.is_running:
+            try:
+                data = self._device.read(64)
+                if data:
+                    try:
+                        self._state = self.parse_data(data)
+                        if self._debug:
+                            print(self._state)
+                    except Exception as e:
+                        print(f"Error parsing data: {e}")
+            except Exception as e:
+                print(f"Error reading from device: {e}")
+
+        self._device.close()
+
+    def stop(self):
+        """
+        Stop reading gamepad input.
+        """
+        self.is_running = False
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Test the Gamepad connection", add_help=True
+    )
+    args = parser.parse_args()
+
+    gamepad = Gamepad(debug=True)
+    while True:
+        time.sleep(1.0)

genesis_forge/gamepads/logitech.py

@@ -0,0 +1,63 @@
+"""Logitech F310/F710 Gamepad configuration."""
+
+from .config import Button, DPad, GamepadConfig
+
+VENDOR_ID = 0x046D
+
+LOGITECH_F710_CONFIG: GamepadConfig = {
+    "name": "F710",
+    "vendor_id": VENDOR_ID,
+    "product_id": 0xC219,
+    "mapping": [
+        {"axis": 0, "data": 1},
+        {"axis": 1, "data": 2},
+        {"axis": 2, "data": 3},
+        {"axis": 3, "data": 4},
+        {"dpad": DPad.UP, "data": 5, "bitmask": 15, "value": 0},
+        {"dpad": DPad.DOWN, "data": 5, "bitmask": 15, "value": 4},
+        {"dpad": DPad.RIGHT, "data": 5, "bitmask": 15, "value": 2},
+        {"dpad": DPad.LEFT, "data": 5, "bitmask": 15, "value": 6},
+        {"button": Button.A, "data": 5, "bitmask": 32},
+        {"button": Button.B, "data": 5, "bitmask": 64},
+        {"button": Button.X, "data": 5, "bitmask": 16},
+        {"button": Button.Y, "data": 5, "bitmask": 128},
+        {"button": Button.LB, "data": 6, "bitmask": 1},
+        {"button": Button.RB, "data": 6, "bitmask": 2},
+        {"button": Button.LT, "data": 6, "bitmask": 4},
+        {"button": Button.RT, "data": 6, "bitmask": 8},
+        {"button": Button.BACK, "data": 6, "bitmask": 16},
+        {"button": Button.START, "data": 6, "bitmask": 32},
+        {"button": Button.MODE, "data": 7, "bitmask": 8},
+        {"button": Button.LEFT_JOYSTICK, "data": 6, "bitmask": 64},
+        {"button": Button.RIGHT_JOYSTICK, "data": 6, "bitmask": 128},
+    ],
+}
+
+LOGITECH_F310_CONFIG: GamepadConfig = {
+    "name": "F310",
+    "vendor_id": VENDOR_ID,
+    "product_id": 0xC216,
+    "mapping": [
+        {"axis": 0, "data": 0},
+        {"axis": 1, "data": 1},
+        {"axis": 2, "data": 2},
+        {"axis": 3, "data": 3},
+        {"dpad": DPad.UP, "data": 4, "bitmask": 15, "value": 0},
+        {"dpad": DPad.DOWN, "data": 4, "bitmask": 15, "value": 4},
+        {"dpad": DPad.RIGHT, "data": 4, "bitmask": 15, "value": 2},
+        {"dpad": DPad.LEFT, "data": 4, "bitmask": 15, "value": 6},
+        {"button": Button.A, "data": 4, "bitmask": 32},
+        {"button": Button.B, "data": 4, "bitmask": 64},
+        {"button": Button.X, "data": 4, "bitmask": 16},
+        {"button": Button.Y, "data": 4, "bitmask": 128},
+        {"button": Button.LB, "data": 5, "bitmask": 1},
+        {"button": Button.RB, "data": 5, "bitmask": 2},
+        {"button": Button.LT, "data": 5, "bitmask": 4},
+        {"button": Button.RT, "data": 5, "bitmask": 8},
+        {"button": Button.BACK, "data": 5, "bitmask": 16},
+        {"button": Button.START, "data": 5, "bitmask": 32},
+        {"button": Button.MODE, "data": 6, "bitmask": 8},
+        {"button": Button.LEFT_JOYSTICK, "data": 5, "bitmask": 64},
+        {"button": Button.RIGHT_JOYSTICK, "data": 5, "bitmask": 128},
+    ],
+}

genesis_forge/genesis_env.py

@@ -0,0 +1,218 @@
+"""
+Simplified Spider Robot Environment with Curriculum Learning
+Focuses on core objectives with progressive difficulty
+"""
+
+import math
+import torch
+import genesis as gs
+from gymnasium import spaces
+from genesis.engine.entities import RigidEntity
+from typing import Any, Literal
+
+EnvMode = Literal["train", "eval", "play"]
+
+
+class GenesisEnv:
+    """
+    Base vectorized environment for Genesis.
+
+    Args:
+        num_envs: Number of parallel environments.
+        dt: Simulation time step.
+        max_episode_length_sec: Maximum episode length in seconds.
+        max_episode_random_scaling: Scale the maximum episode length by this amount (+/-) so that not all environments reset at the same time.
+        headless: Whether to run the environment in headless mode.
+        extras_logging_key: The key used, in info/extras dict, which is returned by step and reset functions, to send data to tensorboard by the RL agent.
+    """
+
+    action_space: spaces.Space | None = None
+    observation_space: spaces.Space | None = None
+
+    def __init__(
+        self,
+        num_envs: int = 1,
+        dt: float = 1 / 100,
+        max_episode_length_sec: int | None = 10,
+        max_episode_random_scaling: float = 0.0,
+        headless: bool = True,
+        mode: EnvMode = "train",
+        extras_logging_key: str = "episode",
+    ):
+        self.dt = dt
+        self.device = gs.device
+        self.num_envs = num_envs
+        self.headless = headless
+        self.mode = mode
+        self.scene: gs.Scene = None
+        self.robot: RigidEntity = None
+        self.terrain: RigidEntity = None
+
+        self.extras_logging_key = extras_logging_key
+        self._extras = {}
+        self._extras[extras_logging_key] = {}
+
+        self._actions: torch.Tensor = None
+        self.last_actions: torch.Tensor = None
+
+        self.step_count: int = 0
+        self.episode_length = torch.zeros(
+            (self.num_envs,), device=gs.device, dtype=torch.int32
+        )
+        self.max_episode_length: torch.Tensor = None
+
+        self._max_episode_length_sec = 0.0
+        self._max_episode_random_scaling = 0.0
+        self._base_max_episode_length = None
+        if max_episode_length_sec and max_episode_length_sec > 0:
+            self._max_episode_random_scaling = max_episode_random_scaling / self.dt
+            self.max_episode_length = torch.zeros(
+                (self.num_envs,), device=gs.device, dtype=gs.tc_int
+            )
+            self.max_episode_length[:] = self.set_max_episode_length(
+                max_episode_length_sec
+            )
+
+    """
+    Properties
+    """
+
+    @property
+    def unwrapped(self):
+        """Returns the base non-wrapped environment.
+
+        Returns:
+            Env: The base non-wrapped :class:`GenesisEnv` instance
+        """
+        return self
+
+    @property
+    def max_episode_length_sec(self) -> int | None:
+        """The max episode length, in seconds, for each environment."""
+        return self._max_episode_length_sec
+
+    @property
+    def extras(self) -> dict:
+        """
+        The extras/infos dictionary that should be returned by the step and reset functions.
+        This dictionary will be cleared at the start of every step.
+        """
+        return self._extras
+
+    @property
+    def actions(self) -> torch.Tensor:
+        """The current actions for each environment for this step."""
+        return self._actions
+
+    @actions.setter
+    def actions(self, actions: torch.Tensor):
+        """Set the actions for each environment for this step."""
+        self._actions = actions
+
+    """
+    Utilities
+    """
+
+    def set_max_episode_length(self, max_episode_length_sec: int) -> int:
+        """
+        Set or change the maximum episode length.
+
+        Args:
+            max_episode_length_sec: The maximum episode length in seconds.
+
+        Returns:
+            The maximum episode length in steps.
+        """
+        self._max_episode_length_sec = max_episode_length_sec
+        self._base_max_episode_length = math.ceil(max_episode_length_sec / self.dt)
+        return self._base_max_episode_length
+
+    """
+    Operations
+    """
+
+    def build(self) -> None:
+        """
+        Builds the scene and other supporting components necessary for the training environment.
+        This assumes that the scene has already been constructed and assigned to the <env>.scene attribute.
+        """
+        assert (
+            self.scene is not None
+        ), "The scene must be constructed and assigned to the <env>.scene attribute before building."
+        self.scene.build(n_envs=self.num_envs)
+
+    def step(
+        self, actions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
+        """
+        Take an action for each parallel environment.
+
+        Args:
+            actions: Batch of actions with the :attr:`action_space` shape.
+
+        Returns:
+            Batch of (observations, rewards, terminations, truncations, infos/extras)
+        """
+        self._extras = {}
+        self._extras[self.extras_logging_key] = {}
+        self.step_count += 1
+        self.episode_length += 1
+
+        self.last_actions[:] = self.actions[:]
+        self._actions = actions
+
+        return None, None, None, None, self._extras
+
+    def reset(
+        self,
+        envs_idx: list[int] = None,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        """
+        Reset one or all parallel environments.
+
+        Args:
+            envs_idx: The environment ids to reset. If None, all environments are reset.
+
+        Returns:
+            A batch of observations and info from the vectorized environment.
+        """
+        if envs_idx is None:
+            envs_idx = torch.arange(self.num_envs, device=gs.device)
+
+        # Initial reset, set buffers
+        if self.step_count == 0:
+            self.actions = torch.zeros(
+                (self.num_envs, self.action_space.shape[0]),
+                device=gs.device,
+                dtype=gs.tc_float,
+            )
+            self.last_actions = torch.zeros_like(self.actions, device=gs.device)
+
+        # Actions
+        if envs_idx.numel() > 0:
+            self.actions[envs_idx] = 0.0
+            self.last_actions[envs_idx] = 0.0
+
+            # Episode length
+            self.episode_length[envs_idx] = 0
+
+        # Randomize max episode length for env_ids
+        if (
+            len(envs_idx) > 0
+            and self._max_episode_random_scaling > 0.0
+            and self._base_max_episode_length
+        ):
+            scale = torch.rand((envs_idx.numel(),)) * self._max_episode_random_scaling
+            self.max_episode_length[envs_idx] = torch.round(
+                self._base_max_episode_length + scale
+            ).to(gs.tc_int)
+
+        return None, self.extras
+
+    def close(self):
+        """Close the environment."""
+        pass
+
+    def render(self):
+        """Not implemented."""
+        pass