Cleanups

actsafe/la_mbda/la_mbda.py → actsafe/actsafe/actsafe.py

@@ -8,13 +8,13 @@
 from omegaconf import DictConfig
 
 from actsafe.common.learner import Learner
-from actsafe.la_mbda import rssm
-from actsafe.la_mbda.exploration import make_exploration
-from actsafe.la_mbda.make_actor_critic import make_actor_critic
-from actsafe.la_mbda.multi_reward import MultiRewardBridge
-from actsafe.la_mbda.replay_buffer import ReplayBuffer
-from actsafe.la_mbda.sentiment import make_sentiment
-from actsafe.la_mbda.world_model import WorldModel, evaluate_model, variational_step
+from actsafe.actsafe import rssm
+from actsafe.actsafe.exploration import UniformExploration, make_exploration
+from actsafe.actsafe.make_actor_critic import make_actor_critic
+from actsafe.actsafe.multi_reward import MultiRewardBridge
+from actsafe.actsafe.replay_buffer import ReplayBuffer
+from actsafe.actsafe.sentiment import make_sentiment
+from actsafe.actsafe.world_model import WorldModel, evaluate_model, variational_step
 from actsafe.rl.epoch_summary import EpochSummary
 from actsafe.rl.metrics import MetricsMonitor
 from actsafe.rl.trajectory import TrajectoryData, Transition
@@ -53,7 +53,7 @@ def init(cls, batch_size: int, cell: rssm.RSSM, action_dim: int) -> "AgentState"
         return self
 
 
-class LaMBDA:
+class ActSafe:
     def __init__(
         self,
         observation_space: Box,
@@ -99,6 +99,7 @@ def __init__(
             action_dim,
             next(self.prng),
         )
+        self.offline = UniformExploration(action_dim)
         self.state = AgentState.init(
             config.training.parallel_envs, self.model.cell, action_dim
         )
@@ -112,6 +113,9 @@ def __init__(
         self.should_explore = Until(
             config.agent.exploration_steps, environment_steps_per_agent_step
         )
+        self.should_collect_offline = Until(
+            config.agent.offline_steps, environment_steps_per_agent_step
+        )
         learn_model_steps = (
             config.agent.learn_model_steps
             if config.agent.learn_model_steps is not None
@@ -128,12 +132,16 @@ def __call__(
     ) -> FloatArray:
         if train and self.should_train() and not self.replay_buffer.empty:
             self.update()
-        policy_fn = (
-            self.exploration.get_policy()
-            if self.should_explore()
-            else self.actor_critic.actor.act
-        )
+        if self.should_collect_offline():
+            policy_fn = self.offline.get_policy()
+        else:
+            policy_fn = (
+                self.exploration.get_policy()
+                if self.should_explore()
+                else self.actor_critic.actor.act
+            )
         self.should_explore.tick()
+        self.should_collect_offline.tick()
         self.learn_model.tick()
         actions, self.state = policy(
             policy_fn,

actsafe/la_mbda/augmented_lagrangian.py → actsafe/actsafe/augmented_lagrangian.py

@@ -6,8 +6,8 @@
 import jax.numpy as jnp
 from jaxtyping import PyTree
 
-from actsafe.la_mbda.actor_critic import ContinuousActor
-from actsafe.la_mbda.safe_actor_critic import ActorEvaluation
+from actsafe.actsafe.actor_critic import ContinuousActor
+from actsafe.actsafe.safe_actor_critic import ActorEvaluation
 
 
 class AugmentedLagrangianUpdate(NamedTuple):

actsafe/la_mbda/dummy_penalizer.py → actsafe/actsafe/dummy_penalizer.py

@@ -4,8 +4,8 @@
 import jax
 from jaxtyping import PyTree
 
-from actsafe.la_mbda.actor_critic import ContinuousActor
-from actsafe.la_mbda.safe_actor_critic import ActorEvaluation
+from actsafe.actsafe.actor_critic import ContinuousActor
+from actsafe.actsafe.safe_actor_critic import ActorEvaluation
 
 
 class DummyPenalizer:

actsafe/la_mbda/exploration.py → actsafe/actsafe/exploration.py

@@ -1,9 +1,9 @@
 import jax
 from omegaconf import DictConfig
 
-from actsafe.la_mbda.opax_bridge import OpaxBridge
-from actsafe.la_mbda.make_actor_critic import make_actor_critic
-from actsafe.la_mbda.sentiment import identity, make_sentiment
+from actsafe.actsafe.opax_bridge import OpaxBridge
+from actsafe.actsafe.make_actor_critic import make_actor_critic
+from actsafe.actsafe.sentiment import identity, make_sentiment
 from actsafe.rl.types import Model, Policy
 
 

actsafe/la_mbda/lbsgd.py → actsafe/actsafe/lbsgd.py

@@ -9,8 +9,8 @@
 
 from actsafe.common.mixed_precision import apply_dtype
 from actsafe.common.pytree_utils import pytrees_unstack
-from actsafe.la_mbda.actor_critic import ContinuousActor
-from actsafe.la_mbda.safe_actor_critic import ActorEvaluation
+from actsafe.actsafe.actor_critic import ContinuousActor
+from actsafe.actsafe.safe_actor_critic import ActorEvaluation
 
 _EPS = 1e-8
 

actsafe/la_mbda/make_actor_critic.py → actsafe/actsafe/make_actor_critic.py

@@ -1,10 +1,10 @@
 import logging
 import numpy as np
-from actsafe.la_mbda.augmented_lagrangian import AugmentedLagrangianPenalizer
-from actsafe.la_mbda.dummy_penalizer import DummyPenalizer
-from actsafe.la_mbda.lbsgd import LBSGDPenalizer
-from actsafe.la_mbda.safe_actor_critic import SafeModelBasedActorCritic
-from actsafe.la_mbda.sentiment import bayes
+from actsafe.actsafe.augmented_lagrangian import AugmentedLagrangianPenalizer
+from actsafe.actsafe.dummy_penalizer import DummyPenalizer
+from actsafe.actsafe.lbsgd import LBSGDPenalizer
+from actsafe.actsafe.safe_actor_critic import SafeModelBasedActorCritic
+from actsafe.actsafe.sentiment import bayes
 
 
 _LOG = logging.getLogger(__name__)

actsafe/la_mbda/multi_reward.py → actsafe/actsafe/multi_reward.py

@@ -1,7 +1,7 @@
 import jax
 import equinox as eqx
-from actsafe.la_mbda.rssm import ShiftScale, State
-from actsafe.la_mbda.world_model import WorldModel
+from actsafe.actsafe.rssm import ShiftScale, State
+from actsafe.actsafe.world_model import WorldModel
 from actsafe.rl.types import Policy, Prediction
 
 

actsafe/la_mbda/opax_bridge.py → actsafe/actsafe/opax_bridge.py

@@ -1,8 +1,8 @@
 import jax
 import equinox as eqx
 from actsafe import opax
-from actsafe.la_mbda.rssm import ShiftScale, State
-from actsafe.la_mbda.world_model import WorldModel
+from actsafe.actsafe.rssm import ShiftScale, State
+from actsafe.actsafe.world_model import WorldModel
 from actsafe.rl.types import Policy, Prediction
 
 

actsafe/la_mbda/safe_actor_critic.py → actsafe/actsafe/safe_actor_critic.py

@@ -9,9 +9,9 @@
 
 from actsafe.common.learner import Learner
 from actsafe.common.mixed_precision import apply_mixed_precision
-from actsafe.la_mbda.rssm import ShiftScale
-from actsafe.la_mbda.sentiment import Sentiment
-from actsafe.la_mbda.actor_critic import ContinuousActor, Critic, actor_entropy
+from actsafe.actsafe.rssm import ShiftScale
+from actsafe.actsafe.sentiment import Sentiment
+from actsafe.actsafe.actor_critic import ContinuousActor, Critic, actor_entropy
 from actsafe.opax import normalized_epistemic_uncertainty
 from actsafe.rl.types import Model, RolloutFn
 from actsafe.rl.utils import nest_vmap

actsafe/la_mbda/sentiment.py → actsafe/actsafe/sentiment.py

@@ -1,7 +1,7 @@
 from typing import Protocol
 import jax
 
-from actsafe.la_mbda.rssm import ShiftScale
+from actsafe.actsafe.rssm import ShiftScale
 from actsafe.opax import normalized_epistemic_uncertainty
 
 

actsafe/la_mbda/world_model.py → actsafe/actsafe/world_model.py

@@ -8,9 +8,9 @@
 
 from actsafe.common.learner import Learner
 from actsafe.common.mixed_precision import apply_mixed_precision
-from actsafe.la_mbda.rssm import RSSM, Features, ShiftScale, State
+from actsafe.actsafe.rssm import RSSM, Features, ShiftScale, State
 from actsafe.rl.types import Prediction
-from actsafe.la_mbda.utils import marginalize_prediction
+from actsafe.actsafe.utils import marginalize_prediction
 from actsafe.rl.types import Policy
 from actsafe.rl.utils import nest_vmap
 

actsafe/configs/agent/la_mbda.yaml → actsafe/configs/agent/actsafe.yaml

@@ -1,7 +1,7 @@
 defaults:
   - penalizer: lbsgd
 
-name: lambda
+name: actsafe
 replay_buffer:
   batch_size: 16
   sequence_length: 50
@@ -52,6 +52,7 @@ safety_slack: 0.
 evaluate_model: false
 exploration_strategy: uniform
 exploration_steps: 5000
+offline_steps: 200000
 learn_model_steps: null
 exploration_reward_scale: 10.0
 exploration_epistemic_scale: 1.

actsafe/configs/config.yaml

@@ -1,6 +1,6 @@
 defaults:
   - _self_
-  - agent: la_mbda
+  - agent: actsafe
   - environment: safe_adaptation_gym
 
 hydra:

actsafe/opax.py

@@ -1,6 +1,6 @@
 import jax
 import jax.numpy as jnp
-from actsafe.la_mbda.rssm import ShiftScale
+from actsafe.actsafe.rssm import ShiftScale
 from actsafe.rl.types import Prediction
 
 _EPS = 1e-5

actsafe/rl/trainer.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 from actsafe import benchmark_suites
-from actsafe.la_mbda.la_mbda import LaMBDA
+from actsafe.actsafe.actsafe import ActSafe
 from actsafe.rl import acting, episodic_async_env
 from actsafe.rl.epoch_summary import EpochSummary
 from actsafe.rl.logging import StateWriter, TrainingLogger
@@ -57,7 +57,7 @@ def __init__(
         self,
         config: DictConfig,
         make_env: EnvironmentFactory,
-        agent: LaMBDA | None = None,
+        agent: ActSafe | None = None,
         start_epoch: int = 0,
         step: int = 0,
         seeds: PRNGSequence | None = None,
@@ -88,10 +88,10 @@ def __enter__(self):
             self.agent = self.make_agent()
         return self
 
-    def make_agent(self) -> LaMBDA:
+    def make_agent(self) -> ActSafe:
         assert self.env is not None
-        if self.config.agent.name == "lambda":
-            agent = LaMBDA(
+        if self.config.agent.name == "actsafe":
+            agent = ActSafe(
                 self.env.observation_space,
                 self.env.action_space,
                 self.config,
@@ -198,7 +198,7 @@ def __init__(
         self,
         config: DictConfig,
         make_env: EnvironmentFactory,
-        agent: LaMBDA | None = None,
+        agent: ActSafe | None = None,
         start_epoch: int = 0,
         step: int = 0,
         seeds: PRNGSequence | None = None,