Gocarx

.github/workflows/ci-build.yml

@@ -4,9 +4,9 @@ name: CI
 
 on:
   push:
-    branches: [ "main", "staging" ]
+    branches: [ "main", "gocarx" ]
   pull_request:
-    branches: [ "main", "staging" ]
+    branches: [ "main", "gocarx" ]
 
 permissions:
   contents: read
@@ -36,5 +36,5 @@ jobs:
         pip install -e .
     - name: Test with pytest
       run: |
-        pytest
+        pytest --capture=no -v waymax
 

.gitignore

@@ -0,0 +1,9 @@
+dataset/training/training_tfexample.tfrecord-00000-of-01000
+dataset/training/training_tfexample.tfrecord-00001-of-01000
+/.vscode
+docs/
+wandb/
+logs/
+out/
+__pycache__
+*.egg-info

Dockerfile

@@ -0,0 +1,3 @@
+FROM ghcr.io/nvidia/jax:jax
+
+CMD ["bash"]

Makefile

@@ -0,0 +1,21 @@
+
+cover_packages=waymax
+
+out=out
+tr=$(out)/test-results
+
+junit=--junitxml=$(tr)/junit.xml
+parallel=-n auto --dist=loadfile
+extra=--capture=no -v
+
+clean-test:
+	poetry run coverage erase
+	rm -rf $(tr) $(tr)
+
+test: clean-test
+	mkdir -p  $(tr)
+	poetry run pytest $(extra) $(junit) waymax
+
+test-parallel: clean-test
+	mkdir -p  $(tr)
+	poetry run pytest $(extra) $(junit) $(parallel) waymax

setup.py

@@ -13,41 +13,44 @@
 # limitations under the License.
 
 """setup.py file for Waymax."""
-from setuptools import find_packages
-from setuptools import setup
-
+from setuptools import find_packages, setup
 
 __version__ = '0.1.0'
 
-
 with open('README.md', encoding='utf-8') as f:
-  _long_description = f.read()
+    _long_description = f.read()
 
 setup(
-    name='waymo-waymax',
-    version=__version__,
-    description='Waymo simulator for autonomous driving',
-    long_description=_long_description,
-    long_description_content_type='text/markdown',
-    author='Waymax team',
-    author_email='[email protected]',
-    python_requires='>=3.10',
-    packages=find_packages(),
-    install_requires=[
-        'numpy>=1.20',
-        'jax>=0.4.6',
-        'tensorflow>=2.11.0',
-        'chex>=0.1.6',
-        'dm_env>=1.6',
-        'flax>=0.6.7',
-        'matplotlib>=3.7.1',
-        'dm-tree>=0.1.8',
-        'immutabledict>=2.2.3',
-        'Pillow>=9.4.0',
-        'mediapy>=1.1.6',
-        'tqdm>=4.65.0',
-        'absl-py>=1.4.0',
-    ],
-    url='https://github.com/waymo-research/waymax',
-    license='Apache-2.0',
+        name='waymax',
+        version=__version__,
+        description='Waymo simulator for autonomous driving',
+        long_description=_long_description,
+        long_description_content_type='text/markdown',
+        author='Waymax team',
+        author_email='[email protected]',
+        python_requires='>=3.10',
+        packages=find_packages(),
+        install_requires=[
+            'numpy>=1.20',
+            'jax>=0.4.6',
+            'jaxtyping',
+            'chex>=0.1.6',
+            'distrax>=0.1.5',
+            'tf-keras', # needed for distrax
+            'dm_env>=1.6',
+            'flax>=0.6.7',
+            'matplotlib<3.10',
+            'dm-tree>=0.1.8',
+            'immutabledict>=2.2.3',
+            'Pillow>=9.4.0',
+            'mediapy>=1.1.6',
+            'moviepy',
+            'imageio',
+            'tqdm>=4.65.0',
+            'absl-py>=1.4.0',
+            'pytest>=6.2.4',
+            "beartype"
+        ],
+        url='https://github.com/waymo-research/waymax',
+        license='Apache-2.0',
 )

waymax/config.py

@@ -15,7 +15,9 @@
 """Configs for Waymax Environments."""
 import dataclasses
 import enum
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Callable
+
+import jax
 
 
 class CoordinateFrame(enum.Enum):
@@ -136,9 +138,25 @@ class LinearCombinationRewardConfig:
     rewards: Dictionary of metric names to floats indicating the weight of each
       metric to create a reward of a linear combination.
   """
-
   rewards: dict[str, float]
 
+  @classmethod
+  def default_gokart(cls) -> 'LinearCombinationRewardConfig':
+    return cls(
+        rewards={"gokart_offroad":-4.0, "gokart_progress": 1.0},
+    )
+
+@dataclasses.dataclass(frozen=True)
+class LinearTransformedRewardConfig(LinearCombinationRewardConfig):
+  """Config listing all metrics and their corresponding transform.
+
+  Attributes:
+    rewards: Dictionary of metric names to floats indicating the weight of each
+      metric to create a reward of a linear combination.
+    transform: Dictionary of metric names to functions that apply an additional transform to the metric
+  """
+  transform: dict[str, Callable[[jax.Array], jax.Array]]
+
 
 class ObjectType(enum.Enum):
   """Types of objects that can be controlled by Waymax."""

waymax/datatypes/__init__.py

@@ -16,6 +16,8 @@
 
 from waymax.datatypes.action import Action
 from waymax.datatypes.action import TrajectoryUpdate
+from waymax.datatypes.action import GoKartTrajectoryUpdate
+from waymax.datatypes.action import GokartAction
 from waymax.datatypes.array import MaskedArray
 from waymax.datatypes.array import PyTree
 from waymax.datatypes.constant import TIME_INTERVAL
@@ -24,6 +26,7 @@
 from waymax.datatypes.object_state import ObjectMetadata
 from waymax.datatypes.object_state import ObjectTypeIds
 from waymax.datatypes.object_state import Trajectory
+from waymax.datatypes.object_state import GokartTrajectory
 from waymax.datatypes.observation import ObjectPose2D
 from waymax.datatypes.observation import Observation
 from waymax.datatypes.observation import observation_from_state
@@ -46,8 +49,10 @@
 from waymax.datatypes.roadgraph import MapElementIds
 from waymax.datatypes.roadgraph import RoadgraphPoints
 from waymax.datatypes.route import Paths
+from waymax.datatypes.route import GoKartPaths
 from waymax.datatypes.simulator_state import get_control_mask
 from waymax.datatypes.simulator_state import SimulatorState
+from waymax.datatypes.simulator_state import GoKartSimState
 from waymax.datatypes.simulator_state import update_state_by_log
 from waymax.datatypes.traffic_lights import TrafficLights
 from waymax.datatypes.traffic_lights import TrafficLightStates

waymax/datatypes/action.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 
 """Dataclass definitions for dynamics models."""
-from typing import Any
+from typing import Any, Sequence
 
 import chex
 import jax
 import jax.numpy as jnp
 
 from waymax.datatypes import operations
+from waymax.utils.classproperty import classproperty
 
 
 # TODO: make Actions inherit from datatypes.MaskedArray.
@@ -112,4 +113,135 @@ def as_action(self) -> Action:
     )
     return Action(data=action, valid=self.valid)
 
+@chex.dataclass
+class GoKartTrajectoryUpdate(TrajectoryUpdate):
+  yaw_rate: jax.Array  # (..., num_objects, 1)
+  acc_x: jax.Array  # (..., num_objects, 1)
+  acc_y: jax.Array  # (..., num_objects, 1)
+
+  def validate(self) -> None:
+    """Validates shape and type."""
+    # Verifies that each element has the same dimensions.
+    chex.assert_equal_shape(
+        [self.x, self.y, self.yaw, self.vel_x, self.vel_y, self.yaw_rate, self.acc_x, self.acc_y, self.valid],
+    )
+    chex.assert_type(
+        [self.x, self.y, self.yaw, self.vel_x, self.vel_y, self.yaw_rate, self.acc_x, self.acc_y, self.valid],
+        [
+            jnp.float32,
+            jnp.float32,
+            jnp.float32,
+            jnp.float32,
+            jnp.float32,
+            jnp.float32,
+            jnp.float32,
+            jnp.float32,
+            jnp.bool_,
+        ],
+    )
+
+  def as_action(self) -> Action:
+    """Returns this trajectory update as a 8D Action for StateDynamics.
+
+    Returns:
+      An action data structure with data of shape (..., 8) containing
+      x, y, yaw, vel_x, vel_y, yaw_rate, acc_x and acc_y.
+    """
+    action = jnp.concatenate(
+        [self.x, self.y, self.yaw, self.vel_x, self.vel_y, self.yaw_rate, self.acc_x, self.acc_y], axis=-1
+    )
+    return Action(data=action, valid=self.valid)
+
+
+@chex.dataclass
+class GokartAction:
+    """
+    Data structure representing the action history of a gokart.
+
+    Attributes:
+
+      steering_angle: The steering angle of the gokart column data type float32.
+      acc_left: The left acceleration of the left wheel of the gokart at each time step of data type float32.
+      acc_right: The right acceleration of the left wheel of the gokart at each time step of data type float32.
+    """
 
+    steering_angle: jax.Array # (..., num_objects, 1)
+    acc_left: jax.Array # (..., num_objects, 1)
+    acc_right: jax.Array # (..., num_objects, 1)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """The Array shape of this history."""
+        return self.steering_angle.shape
+
+    @property
+    def num_actions(self) -> int:
+        """The number of actions."""
+        return self.shape[-2]
+
+    @property
+    def num_timesteps(self) -> int:
+        """The length of this history."""
+        return self.shape[-1]
+
+    @property
+    def acc_leftright(self) -> jax.Array:
+        """Stacked AB action"""
+        return jnp.stack([self.acc_left, self.acc_right], axis=-1)
+
+    @property
+    def torque_vectoring(self) -> jax.Array:
+        """Stacked Torque Vectoring indirect action (acc_right - acc_left)"""
+        return self.acc_right - self.acc_left
+
+    @classproperty
+    def action_fields(self) -> Sequence[str]:
+        """Returns the action fields."""
+        #todo there are better ways to do this
+        return ["steering_angle", "acc_left", "acc_right"]
+
+    @classmethod
+    def zeros(cls, shape: Sequence[int]) -> "GokartAction":
+        """Creates a Trajectory containing zeros of the specified shape."""
+        return cls(
+            steering_angle=jnp.zeros(shape, jnp.float32),
+            acc_left=jnp.zeros(shape, jnp.float32),
+            acc_right=jnp.zeros(shape, jnp.float32),
+        )
+
+    def __eq__(self, other: Any) -> bool:
+        return operations.compare_all_leaf_nodes(self, other)
+
+    def stack_fields(self, field_names: Sequence[str]) -> jax.Array:
+        """Returns a concatenated version of a set of field names for Trajectory."""
+        return jnp.stack([getattr(self, field_name) for field_name in field_names], axis=-1)
+
+    def set_actions(self, action: Action, timestep: jax.typing.ArrayLike) -> "GokartAction":
+        """Return a new action history updated at timestep with the new action."""
+        return self.replace(
+            steering_angle=self.steering_angle.at[..., timestep].set(action.data[0]),
+            acc_left=self.acc_left.at[..., timestep].set(action.data[1]),
+            acc_right=self.acc_right.at[..., timestep].set(action.data[2]),
+        )
+
+    def validate(self):
+        """Validates shape and type."""
+        chex.assert_equal_shape(
+            [
+                self.steering_angle,
+                self.acc_left,
+                self.acc_right,
+            ]
+        )
+        chex.assert_type(
+            [
+                self.steering_angle,
+                self.acc_left,
+                self.acc_right,
+            ],
+            [
+                jnp.float32,
+                jnp.float32,
+                jnp.float32,
+            ],
+        )