Sgym unsupervised (#33)

* Update SAG

* No need to really reset tasks for now

* Improve config file

* Add Bhavi scale

* Back to circle, use Bhavi scale

* Update hparams

* Set default epistemic scale to 1.

* Fix tests

Author: yardenas

poetry.lock

@@ -54,5 +54,6 @@ exploration_strategy: uniform
 exploration_steps: 5000
 learn_model_steps: null
 exploration_reward_scale: 10.0
+exploration_epistemic_scale: 1.
 unsupervised: false
 reward_scale: 1.

safe_opax/configs/agent/la_mbda.yaml

@@ -31,4 +31,4 @@ agent:
     batch_size: 4
     sequence_length: 16
   exploration_steps: 750
-  unsupervised: true
+  unsupervised: true

safe_opax/configs/experiment/debug_unsupervised.yaml

@@ -4,19 +4,21 @@ defaults:
 
 training:
   trainer: unsupervised
-  epochs: 200
+  epochs: 100
   safe: true
   action_repeat: 2
-  episodes_per_epoch: 5
   exploration_steps: 1000000
   train_task_name: unsupervised
   test_task_name: go_to_goal
 
 environment:
   safe_adaptation_gym:
-    robot_name: doggo
+    robot_name: point
 
 agent:
   exploration_strategy: opax
   exploration_steps: 1000000
-  unsupervised: true
+  unsupervised: true
+  learn_model_steps: 1000000
+  exploration_epistemic_scale: 15.0
+  exploration_reward_scale: 25.0

safe_opax/configs/experiment/unsupervised_safety_gym.yaml

@@ -50,14 +50,15 @@ def __init__(
             ),
         )
         self.reward_scale = config.agent.exploration_reward_scale
+        self.epistemic_scale = config.agent.exploration_epistemic_scale
 
     def update(
         self,
         model: Model,
         initial_states: jax.Array,
         key: jax.Array,
     ) -> dict[str, float]:
-        model = OpaxBridge(model, self.reward_scale)
+        model = OpaxBridge(model, self.reward_scale, self.epistemic_scale)
         outs = self.actor_critic.update(model, initial_states, key)
         outs = {f"{_append_opax(k)}": v for k, v in outs.items()}
         return outs
@@ -77,6 +78,5 @@ def __init__(self, action_dim: int):
         self.action_dim = action_dim
         self.policy = lambda _, key: jax.random.uniform(key, (self.action_dim,))
 
-
     def get_policy(self) -> Policy:
         return self.policy

safe_opax/la_mbda/exploration.py

@@ -9,6 +9,7 @@
 class OpaxBridge(eqx.Module):
     model: WorldModel
     reward_scale: float = eqx.field(static=True)
+    reward_epistemic_scale: float = eqx.field(static=True)
 
     def sample(
         self,
@@ -21,4 +22,6 @@ def sample(
             horizon, initial_state, key, policy
         )
         trajectory, distributions = samples
-        return opax.modify_reward(trajectory, distributions, self.reward_scale)
+        return opax.modify_reward(
+            trajectory, distributions, self.reward_scale, self.reward_epistemic_scale
+        )

safe_opax/la_mbda/opax_bridge.py

@@ -10,9 +10,12 @@ def modify_reward(
     trajectory: Prediction,
     distributions: ShiftScale,
     scale: float = 1.0,
+    epistemic_scale: float = 1.0,
     stop_grad: bool = True,
 ) -> tuple[Prediction, ShiftScale]:
-    new_rewards = normalized_epistemic_uncertainty(distributions) * scale
+    new_rewards = (
+        normalized_epistemic_uncertainty(distributions, scale=epistemic_scale) * scale
+    )
     if stop_grad:
         new_rewards = jax.lax.stop_gradient(new_rewards)
     return Prediction(
@@ -23,10 +26,15 @@ def modify_reward(
 
 
 def normalized_epistemic_uncertainty(
-    distributions: ShiftScale, axis: int = 0
+    distributions: ShiftScale, axis: int = 0, scale: float = 1.0
 ) -> jnp.ndarray:
     epistemic_uncertainty = distributions.shift.var(axis)
     aleatoric_uncertainty = (distributions.scale**2).mean(axis)
     return 0.5 * jnp.log(
-        1.0 + (epistemic_uncertainty.mean(-1) / (aleatoric_uncertainty.mean(-1) + _EPS))
+        1.0
+        + (
+            scale
+            * epistemic_uncertainty.mean(-1)
+            / (aleatoric_uncertainty.mean(-1) + _EPS)
+        )
     )

safe_opax/opax.py

@@ -213,6 +213,8 @@ def __init__(
         super().__init__(config, make_env, agent, start_epoch, step, seeds)
         self.test_task_name = self.config.training.test_task_name
         self.train_task_name = self.config.training.train_task_name
+        # After a few iterations, we realized `test_tasks` are not useful, as we just use multiple rewards.
+        # just ignore this.
         self.test_tasks: list[Task] | None = None
 
     def __enter__(self):
@@ -233,7 +235,6 @@ def __enter__(self):
                 get_task(self.test_task_name)
                 for _ in range(self.config.training.parallel_envs)
             ]
-            self.env.reset(options={"task": self.test_tasks})
         return self
 
     def _run_training_epoch(
@@ -250,7 +251,6 @@ def _run_training_epoch(
                 for _ in range(self.config.training.parallel_envs)
             ]
             assert self.env is not None
-            self.env.reset(options={"task": self.test_tasks})
             assert self.agent is not None
         return outs
 

safe_opax/rl/trainer.py

@@ -43,4 +43,4 @@ def test_epoch(trainer):
     with trainer as trainer:
         with patch.object(trainer.env, "reset", wraps=trainer.env.reset) as mock:
             trainer.train(1)
-    assert mock.call_count == 4
+    assert mock.call_count == 3