Domain randomization rccar (#8)

* Add nominal model

* Add bounds to car models

* Add randomization functions

* Sampling seems to work

* Add domain randomization test

* Remove unused function

Author: yardenas

ss2r/benchmark_suites/__init__.py

@@ -16,24 +16,54 @@ def make(cfg):
         return make_rccar_envs(cfg)
 
 
+def prepare_randomization_fn(key, num_envs, cfg, task_name):
+    randomize_fn = lambda sys, rng: randomization_fns[task_name](sys, rng, cfg)
+    v_randomization_fn = functools.partial(
+        randomize_fn, rng=jax.random.split(key, num_envs)
+    )
+    vf_randomization_fn = lambda sys: v_randomization_fn(sys)[:-1]  # type: ignore
+    params_fn = lambda sys: v_randomization_fn(sys)[-1]
+    return vf_randomization_fn, params_fn
+
+
 def make_rccar_envs(cfg):
     task_cfg = dict(get_task_config(cfg))
     task_cfg.pop("domain_name")
+    task_cfg.pop("task_name")
     train_car_params = task_cfg.pop("train_car_params")
     eval_car_params = task_cfg.pop("eval_car_params")
-    train_env = rccar.RCCar(train_car_params, **task_cfg)
+    train_key, eval_key = jax.random.split(jax.random.PRNGKey(cfg.training.seed))
+    train_env = rccar.RCCar(train_car_params["nominal"], **task_cfg)
+    train_randomization_fn, params_fn = (
+        prepare_randomization_fn(
+            train_key,
+            cfg.training.num_envs,
+            train_car_params["bounds"],
+            cfg.environment.task_name,
+        )
+        if cfg.training.train_domain_randomization
+        else (None, None)
+    )
     train_env = envs.training.wrap(
         train_env,
         episode_length=cfg.training.episode_length,
         action_repeat=cfg.training.action_repeat,
+        randomization_fn=train_randomization_fn,
+    )
+    eval_env = rccar.RCCar(eval_car_params["nominal"], **task_cfg)
+    eval_randomization_fn, _ = prepare_randomization_fn(
+        eval_key,
+        cfg.training.num_eval_envs,
+        eval_car_params["bounds"],
+        cfg.environment.task_name,
     )
-    eval_env = rccar.RCCar(eval_car_params, **task_cfg)
     eval_env = envs.training.wrap(
         eval_env,
         episode_length=cfg.training.episode_length,
         action_repeat=cfg.training.action_repeat,
+        randomization_fn=eval_randomization_fn,
     )
-    return train_env, eval_env, None
+    return train_env, eval_env, params_fn
 
 
 def make_brax_envs(cfg):
@@ -43,20 +73,10 @@ def make_brax_envs(cfg):
     )
     eval_env = envs.get_environment(task_cfg.task_name, backend=cfg.environment.backend)
     train_key, eval_key = jax.random.split(jax.random.PRNGKey(cfg.training.seed))
-
-    def prepare_randomization_fn(key, num_envs):
-        randomize_fn = lambda sys, rng: randomization_fns[task_cfg.task_name](
-            sys, rng, task_cfg
-        )
-        v_randomization_fn = functools.partial(
-            randomize_fn, rng=jax.random.split(key, num_envs)
-        )
-        vf_randomization_fn = lambda sys: v_randomization_fn(sys)[:-1]  # type: ignore
-        params_fn = lambda sys: v_randomization_fn(sys)[-1]
-        return vf_randomization_fn, params_fn
-
     train_randomization_fn, params_fn = (
-        prepare_randomization_fn(train_key, cfg.training.num_envs)
+        prepare_randomization_fn(
+            train_key, cfg.training.num_envs, task_cfg, task_cfg.task_name
+        )
         if cfg.training.train_domain_randomization
         else (None, None)
     )
@@ -67,7 +87,7 @@ def prepare_randomization_fn(key, num_envs):
         randomization_fn=train_randomization_fn,
     )
     eval_randomization_fn, _ = prepare_randomization_fn(
-        eval_key, cfg.training.num_eval_envs
+        eval_key, cfg.training.num_eval_envs, task_cfg, task_cfg.task_name
     )
     eval_env = envs.training.wrap(
         eval_env,

ss2r/benchmark_suites/brax/__init__.py

@@ -1,8 +1,10 @@
 from ss2r.benchmark_suites.brax.cartpole import cartpole
+from ss2r.benchmark_suites.rccar import rccar
 
 randomization_fns = {
     "cartpole_swingup": cartpole.domain_randomization,
     "cartpole_swingup_sparse": cartpole.domain_randomization,
     "cartpole_balance": cartpole.domain_randomization,
     "inverted_pendulum": cartpole.domain_randomization,
+    "rccar": rccar.domain_randomization,
 }

ss2r/benchmark_suites/rccar/model.py

@@ -19,7 +19,6 @@ class CarParams:
     c_m_1: max current of motor: [0.2 - 0.5] c_m_2: motor resistance due to shaft: [0.01 - 0.15]
     """
 
-    car_id: int = 2
     m: jax.Array = jnp.array(1.65)  # [0.04, 0.08]
     i_com: jax.Array = jnp.array(2.78e-05)  # [1e-6, 5e-6]
     l_f: jax.Array = jnp.array(0.13)  # [0.025, 0.05]

ss2r/benchmark_suites/rccar/rccar.py

@@ -2,8 +2,11 @@
 from typing import Tuple
 
 import jax
+import jax.flatten_util
 import jax.numpy as jnp
+import jax.tree_util as jtu
 from brax.envs.base import Env, State
+from omegaconf import OmegaConf
 
 from ss2r.benchmark_suites import rewards
 from ss2r.benchmark_suites.rccar.model import CarParams, RaceCarDynamics
@@ -13,6 +16,38 @@
 )
 
 
+def domain_randomization(sys, rng, cfg):
+    def sample_from_bounds(value, key):
+        """
+        Sample from a JAX uniform distribution if the value is a list of two elements.
+        """
+        if isinstance(value, list) and len(value) == 2:
+            lower, upper = value
+            # Sample from jax.random.uniform with the given key
+            return jax.random.uniform(key, shape=(), minval=lower, maxval=upper)
+        return value
+
+    @jax.vmap
+    def randomize(rng):
+        bounds = CarParams(**cfg)
+        # Define a custom tree structure that treats lists as leaves
+        treedef = jtu.tree_structure(bounds, is_leaf=lambda x: isinstance(x, list))
+        # Generate random keys only for the relevant leaves (i.e., lists with 2 elements)
+        keys = jax.random.split(rng, num=treedef.num_leaves)
+        # Rebuild the tree with the keys, only where there are valid leaves
+        keys = jtu.tree_unflatten(treedef, keys)
+        # Map over the tree, generating random values where needed
+        sys = jtu.tree_map(
+            sample_from_bounds, bounds, keys, is_leaf=lambda x: isinstance(x, list)
+        )
+        return sys, jax.flatten_util.ravel_pytree(sys)[0]
+
+    cfg = OmegaConf.to_container(cfg)
+    in_axes = jax.tree_map(lambda _: 0, sys)
+    sys, params = randomize(rng)
+    return sys, in_axes, params
+
+
 def rotate_coordinates(state: jnp.array, encode_angle: bool = False) -> jnp.array:
     x_pos, x_vel = (
         state[..., 0:1],

ss2r/configs/environment/eval_car_params/car_1.yaml

@@ -1,18 +1,39 @@
-use_blend: 0.0
-m: 1.65
-l_f: 0.13
-l_r: 0.17
-angle_offset: 0.0156
-b_f: 2.58
-b_r: 3.39
-blend_ratio_lb: 0.01
-blend_ratio_ub: 0.01
-c_d: 0.41464928
-c_f: 1.2
-c_m_1: 10.701814
-c_m_2: 1.4208151
-c_r: 1.27
-d_f: 0.02
-d_r: 0.017
-i_com: 0.01
-steering_limit: 0.3543
+nominal:
+  use_blend: 0.0
+  m: 1.65
+  l_f: 0.13
+  l_r: 0.17
+  angle_offset: 0.0156
+  b_f: 2.58
+  b_r: 3.39
+  blend_ratio_lb: 0.01
+  blend_ratio_ub: 0.01
+  c_d: 0.41464928
+  c_f: 1.2
+  c_m_1: 10.701814
+  c_m_2: 1.4208151
+  c_r: 1.27
+  d_f: 0.02
+  d_r: 0.017
+  i_com: 0.01
+  steering_limit: 0.3543
+
+bounds:
+  use_blend: [0.0, 0.0]
+  m: [1.6, 1.7]
+  l_f: [0.11, 0.15]
+  l_r: [0.15, 0.19]
+  angle_offset: [0.001, 0.03]
+  b_f: [2.2, 2.8]
+  b_r: [2.0, 6.0]
+  blend_ratio_lb: [0.4, 0.4]
+  blend_ratio_ub: [0.5, 0.5]
+  c_d: [0.3, 0.5]
+  c_f: [1.2, 1.2]
+  c_m_1: [8.0, 13.0]
+  c_m_2: [1.1, 1.7]
+  c_r: [1.27, 1.27]
+  d_f: [0.02, 0.02]
+  d_r: [0.017, 0.017]
+  i_com: [0.01, 0.1]
+  steering_limit: [0.20, 0.5]

ss2r/configs/environment/eval_car_params/car_1_blend.yaml

@@ -1,18 +1,39 @@
-use_blend: 1.0
-m: 1.65
-l_f: 0.13
-l_r: 0.17
-angle_offset: -0.0213
-b_f: 1.8966477
-b_r: 6.2884626
-blend_ratio_lb: 0.06637411
-blend_ratio_ub: 0.00554
-c_d: 0.0
-c_f: 1.5381637
-c_m_1: 11.102413
-c_m_2: 1.3169205
-c_r: 1.186591
-d_f: 0.5968191
-d_r: 0.42716035
-i_com: 0.0685434
-steering_limit: 0.6337473
+nominal:
+  use_blend: 1.0
+  m: 1.65
+  l_f: 0.13
+  l_r: 0.17
+  angle_offset: -0.0213
+  b_f: 1.8966477
+  b_r: 6.2884626
+  blend_ratio_lb: 0.06637411
+  blend_ratio_ub: 0.00554
+  c_d: 0.0
+  c_f: 1.5381637
+  c_m_1: 11.102413
+  c_m_2: 1.3169205
+  c_r: 1.186591
+  d_f: 0.5968191
+  d_r: 0.42716035
+  i_com: 0.0685434
+  steering_limit: 0.6337473
+
+bounds:
+  use_blend: [1.0, 1.0]
+  m: [1.6, 1.7]
+  l_f: [0.125, 0.135]
+  l_r: [0.165, 0.175]
+  angle_offset: [-0.025, 0.025]
+  b_f: [1.3, 3.0]
+  b_r: [4.0, 10.0]
+  blend_ratio_lb: [0.01, 0.1]
+  blend_ratio_ub: [0.0, 0.2]
+  c_d: [0.0, 0.0]
+  c_f: [1.2, 1.8]
+  c_m_1: [10.0, 12.0]
+  c_m_2: [1.1, 1.5]
+  c_r: [0.9, 1.5]
+  d_f: [0.35, 0.65]
+  d_r: [0.3, 0.6]
+  i_com: [0.05, 0.09]
+  steering_limit: [0.5, 0.9]

ss2r/configs/environment/eval_car_params/car_2.yaml

@@ -1,18 +1,39 @@
-use_blend: 0.0
-m: 1.65
-l_f: 0.13
-l_r: 0.17
-angle_offset: 0.0
-b_f: 2.58
-b_r: 5.0
-blend_ratio_lb: 0.01
-blend_ratio_ub: 0.01
-c_d: 0.0
-c_f: 1.2
-c_m_1: 8.0
-c_m_2: 1.5
-c_r: 1.27
-d_f: 0.02
-d_r: 0.017
-i_com: 0.01
-steering_limit: 0.3
+nominal:
+  use_blend: 0.0
+  m: 1.65
+  l_f: 0.13
+  l_r: 0.17
+  angle_offset: 0.0
+  b_f: 2.58
+  b_r: 5.0
+  blend_ratio_lb: 0.01
+  blend_ratio_ub: 0.01
+  c_d: 0.0
+  c_f: 1.2
+  c_m_1: 8.0
+  c_m_2: 1.5
+  c_r: 1.27
+  d_f: 0.02
+  d_r: 0.017
+  i_com: 0.01
+  steering_limit: 0.3
+
+bounds:
+  use_blend: [0.0, 0.0]
+  m: [1.6, 1.7]
+  l_f: [0.11, 0.15]
+  l_r: [0.15, 0.19]
+  angle_offset: [-0.15, 0.15]
+  b_f: [2.4, 2.6]
+  b_r: [2.0, 8.0]
+  blend_ratio_lb: [0.4, 0.4]
+  blend_ratio_ub: [0.5, 0.5]
+  c_d: [0.01, 0.01]
+  c_f: [1.2, 1.2]
+  c_m_1: [6.0, 10.0]
+  c_m_2: [1.0, 1.8]
+  c_r: [1.27, 1.27]
+  d_f: [0.02, 0.02]
+  d_r: [0.017, 0.017]
+  i_com: [0.01, 0.1]
+  steering_limit: [0.15, 0.4]

ss2r/configs/environment/eval_car_params/car_2_blend.yaml

@@ -1,18 +1,39 @@
-use_blend: 1.0
-m: 1.65
-l_f: 0.13
-l_r: 0.17
-angle_offset: 0.0
-b_f: 2.75
-b_r: 5.0
-blend_ratio_lb: 0.001
-blend_ratio_ub: 0.017
-c_d: 0.0
-c_f: 1.45
-c_m_1: 8.2
-c_m_2: 1.25
-c_r: 1.3
-d_f: 0.4
-d_r: 0.3
-i_com: 0.06
-steering_limit: 0.6
+nominal:
+  use_blend: 1.0
+  m: 1.65
+  l_f: 0.13
+  l_r: 0.17
+  angle_offset: 0.0
+  b_f: 2.75
+  b_r: 5.0
+  blend_ratio_lb: 0.001
+  blend_ratio_ub: 0.017
+  c_d: 0.0
+  c_f: 1.45
+  c_m_1: 8.2
+  c_m_2: 1.25
+  c_r: 1.3
+  d_f: 0.4
+  d_r: 0.3
+  i_com: 0.06
+  steering_limit: 0.6
+
+bounds:
+  use_blend: [1.0, 1.0]
+  m: [1.6, 1.7]
+  l_f: [0.125, 0.135]
+  l_r: [0.165, 0.175]
+  angle_offset: [-0.15, 0.15]
+  b_f: [2.0, 4.0]
+  b_r: [3.0, 10.0]
+  blend_ratio_lb: [0.0001, 0.1]
+  blend_ratio_ub: [0.0001, 0.2]
+  c_d: [0.0, 0.0]
+  c_f: [1.1, 2.0]
+  c_m_1: [6.5, 10.0]
+  c_m_2: [1.0, 1.5]
+  c_r: [0.4, 2.0]
+  d_f: [0.25, 0.6]
+  d_r: [0.15, 0.45]
+  i_com: [0.03, 0.18]
+  steering_limit: [0.4, 0.75]