Rssm init (#21)

* Extend initialization tricks

* Add parameters for doggo

* Fix bug with rendering

Author: yardenas

pyproject.toml

@@ -2,7 +2,7 @@
 name = "safe-opax"
 version = "0.1.0"
 description = ""
-authors = ["Yarden <[email protected]>"]
+authors = ["Yarden As"]
 readme = "README.md"
 
 [tool.poetry.dependencies]

safe_opax/benchmark_suites/safe_adaptation_gym/__init__.py

@@ -44,7 +44,7 @@ def make_env():
             task_name=task,
             seed=cfg.training.seed,
             rgb_observation=task_cfg.image_observation.enabled,
-            render_lidar_and_collision=not task_cfg.image_observation.enabled,
+            render_lidar_and_collision=False,
         )
         env = SafeAdaptationEnvCompatibility(env)
         if (

safe_opax/common/pytree_utils.py

@@ -1,4 +1,6 @@
 import jax
+import jax.numpy as jnp
+
 
 def pytrees_unstack(pytree):
     leaves, treedef = jax.tree_flatten(pytree)
@@ -9,3 +11,8 @@ def pytrees_unstack(pytree):
             new_leaves[i].append(leaf[i])
     new_trees = [treedef.unflatten(leaf) for leaf in new_leaves]
     return new_trees
+
+
+def pytrees_stack(pytrees, axis=0):
+    results = jax.tree_map(lambda *values: jnp.stack(values, axis=axis), *pytrees)
+    return results

safe_opax/configs/agent/la_mbda.yaml

@@ -8,7 +8,7 @@ replay_buffer:
   capacity: 1000
 sentiment:
   ensemble_size: 5
-  model_initialization_scale: 0.5
+  model_initialization_scale: null
   constraint_pessimism: null
   objective_optimism: null
 model:

safe_opax/configs/config.yaml

@@ -42,7 +42,7 @@ training:
   episodes_per_epoch: 5
   epochs: 200
   action_repeat: 1
-  render_episodes: 1
+  render_episodes: 0
   parallel_envs: 10
   scale_reward: 1.
   exploration_steps: 5000

safe_opax/configs/experiment/safe_sparse_cartpole.yaml

@@ -14,7 +14,7 @@ training:
 
 agent:
   exploration_strategy: opax
-  exploration_steps: 1000000
+  exploration_steps: 0
   actor:
     init_stddev: 0.025
   sentiment:

safe_opax/configs/experiment/safety_gym_doggo.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+defaults:
+  - override /environment: safe_adaptation_gym
+
+training:
+  epochs: 100
+  safe: true
+  action_repeat: 2
+
+environment:
+  safe_adaptation_gym:
+    robot_name: doggo
+
+agent:
+  exploration_steps: 0
+  actor:
+    initialization_scale: 1.

safe_opax/la_mbda/rssm.py

@@ -1,13 +1,13 @@
-from functools import partial
 from typing import NamedTuple
 
 import distrax as dtx
 import equinox as eqx
 import jax
+import jax.flatten_util
 import jax.nn as jnn
 import jax.numpy as jnp
 
-from safe_opax.rl.utils import glorot_uniform, init_linear_weights
+from safe_opax.rl.utils import init_linear_weights_and_biases
 
 
 class State(NamedTuple):
@@ -127,27 +127,26 @@ def __init__(
         embedding_size: int,
         action_dim: int,
         ensemble_size: int,
-        initialization_scale: float,
+        initialization_scale: float | None = None,
         *,
         key: jax.Array,
     ):
         self.ensemble_size = ensemble_size
         prior_key, posterior_key = jax.random.split(key)
-        make_priors = eqx.filter_vmap(
-            lambda key: init_linear_weights(
-                Prior(
-                    deterministic_size,
-                    stochastic_size,
-                    hidden_size,
-                    action_dim,
-                    key,
-                ),
-                partial(glorot_uniform, scale=initialization_scale),
-                key,
-            )
+        dummy_prior = Prior(
+            deterministic_size,
+            stochastic_size,
+            hidden_size,
+            action_dim,
+            key,
+        )
+        initialization_scale = (
+            initialization_scale
+            if initialization_scale is not None
+            else jax.flatten_util.ravel_pytree(dummy_prior)[0].std()
         )
-        self.priors = make_priors(
-            jnp.asarray(jax.random.split(prior_key, ensemble_size))
+        self.priors = jitter_priors(
+            dummy_prior, prior_key, initialization_scale, ensemble_size
         )
         self.posterior = Posterior(
             deterministic_size,
@@ -214,3 +213,14 @@ def _priors_predict(
         in_axes=(eqx.if_array(0), prev_state_in_axis, action_in_axis),
     )
     return priors_fn(priors, prev_state, action)
+
+
+def jitter_priors(
+    prior: Prior, key: jax.Array, scale: float, ensemble_size: int
+) -> Prior:
+    make_priors = eqx.filter_vmap(
+        lambda key: init_linear_weights_and_biases(
+            prior, lambda x, subkey: x + scale * jax.random.normal(subkey, x.shape), key
+        )
+    )
+    return make_priors(jnp.asarray(jax.random.split(key, ensemble_size)))

safe_opax/rl/utils.py

@@ -84,22 +84,35 @@ def rl_initialize_weights_trick(model, bias_shift=0.0, weight_scale=0.01):
         model.layers[-1].weight * weight_scale,
     )
     model = eqx.tree_at(
-        lambda model: model.layers[-1].bias, model, model.layers[-1].bias * 0. + bias_shift
+        lambda model: model.layers[-1].bias,
+        model,
+        model.layers[-1].bias * 0.0 + bias_shift,
     )
     return model
 
 
-def init_linear_weights(model, init_fn, key):
+def init_linear_weights_and_biases(model, init_fn, key):
     is_linear = lambda x: isinstance(x, eqx.nn.Linear)
     get_weights = lambda m: [
         x.weight
         for x in jax.tree_util.tree_leaves(m, is_leaf=is_linear)
         if is_linear(x)
     ]
+    get_biases = lambda m: [
+        x.bias
+        for x in jax.tree_util.tree_leaves(m, is_leaf=is_linear)
+        if is_linear(x) and x.bias is not None
+    ]
     weights = get_weights(model)
+    biases = get_biases(model)
     new_weights = [
         init_fn(weight, subkey)
         for weight, subkey in zip(weights, jax.random.split(key, len(weights)))
     ]
+    new_biases = [
+        init_fn(bias, subkey)
+        for bias, subkey in zip(biases, jax.random.split(key, len(biases)))
+    ]
     new_model = eqx.tree_at(get_weights, model, new_weights)
+    new_model = eqx.tree_at(get_biases, new_model, new_biases)
     return new_model