Merge branch 'scalar_loss' of github.com:UT-Austin-RPL/amago into scalar_loss

Author: jakegrigsby

amago/__init__.py

@@ -1,4 +1,9 @@
-__version__ = "3.1.2"
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("amago")
+except PackageNotFoundError:
+    __version__ = "unknown"
 
 from .experiment import Experiment
 from .agent import Agent, register_agent, get_agent_cls, list_registered_agents

amago/agent.py

@@ -1234,9 +1234,7 @@ def __init__(
         actor_type: Type[actor_critic.BaseActorHead] = actor_critic.Actor,
         critic_type: Type[actor_critic.BaseCriticHead] = actor_critic.NCriticsTwoHot,
         pass_obs_keys_to_actor: Optional[Iterable[str]] = None,
-        dpg_temperature: float = 1.0,
     ):
-        self.dpg_temperature = dpg_temperature
         super().__init__(
             obs_space=obs_space,
             rl2_space=rl2_space,
@@ -1502,9 +1500,7 @@ def forward(self, batch: Batch, log_step: bool):
             a_agent_dpg = torch.stack([a_dist.rsample() for _ in range(K_a)], dim=0)
             q_s_a_agent = self.maximized_critics(s_rep.detach(), a_agent_dpg)
             q_s_a_agent = self.popart.normalize_values(
-                self.maximized_critics.bin_dist_to_raw_vals(
-                    q_s_a_agent, temperature=self.dpg_temperature
-                ).mean(0).min(2).values
+                self.maximized_critics.bin_dist_to_raw_vals(q_s_a_agent).mean(0).min(2).values
             )
             actor_loss += self.online_coeff * -(q_s_a_agent[:, :-1, ...])
 

amago/nets/actor_critic.py

@@ -600,25 +600,18 @@ def critic_network_forward(
         )
         return pyd.Categorical(logits=outputs)
 
-    def bin_dist_to_raw_vals(
-        self, bin_dist: pyd.Categorical, temperature: float = 1.0
-    ) -> torch.Tensor:
+    def bin_dist_to_raw_vals(self, bin_dist: pyd.Categorical) -> torch.Tensor:
         """Convert a categorical distribution over bins to a scalar.
 
         Args:
             bin_dist: The categorical distribution over bins (output of `forward`).
-            temperature: Temperature for softening the distribution.
 
         Returns:
             The scalar value.
         """
         assert isinstance(bin_dist, pyd.Categorical)
-        if temperature == 1.0:
-            probs = bin_dist.probs
-        else:
-            probs = torch.softmax(bin_dist.logits / temperature, dim=-1)
-        bin_vals = self.bin_vals.to(probs.device, dtype=probs.dtype)
-        exp_val = (probs * bin_vals).sum(-1, keepdims=True)
+        bin_vals = self.bin_vals.to(bin_dist.probs.device, dtype=bin_dist.probs.dtype)
+        exp_val = (bin_dist.probs * bin_vals).sum(-1, keepdims=True)
         return self.invert_bins(exp_val)
 
     def raw_vals_to_labels(self, raw_td_target: torch.Tensor) -> torch.Tensor:

docs/api/amago.agent.rst

@@ -10,11 +10,16 @@ amago.agent
 
       binary_filter
       exp_filter
+      register_agent
+      get_agent_cls
+      list_registered_agents
+      nstep_return
 
    .. rubric:: Classes
 
    .. autosummary::
 
+      BaseAgent
       Agent
       MultiTaskAgent
       Multigammas

docs/api/amago.envs.exploration.rst

@@ -4,11 +4,19 @@ amago.envs.exploration
 .. automodule:: amago.envs.exploration
 
 
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      register_exploration
+      get_exploration_cls
+      list_registered_explorations
+
    .. rubric:: Classes
 
    .. autosummary::
 
-      BilevelEpsilonGreedy
-      EpsilonGreedy
       ExplorationWrapper
+      EpsilonGreedy
+      BilevelEpsilonGreedy
 

docs/api/amago.nets.traj_encoders.rst

@@ -4,13 +4,21 @@ amago.nets.traj\_encoders
 .. automodule:: amago.nets.traj_encoders
 
 
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      register_traj_encoder
+      get_traj_encoder_cls
+      list_registered_traj_encoders
+
    .. rubric:: Classes
 
    .. autosummary::
 
+      TrajEncoder
       FFTrajEncoder
       GRUTrajEncoder
-      MambaTrajEncoder
       TformerTrajEncoder
-      TrajEncoder
+      MambaTrajEncoder
 

docs/api/amago.nets.tstep_encoders.rst

@@ -4,11 +4,19 @@ amago.nets.tstep\_encoders
 .. automodule:: amago.nets.tstep_encoders
 
 
+   .. rubric:: Functions
+
+   .. autosummary::
+
+      register_tstep_encoder
+      get_tstep_encoder_cls
+      list_registered_tstep_encoders
+
    .. rubric:: Classes
 
    .. autosummary::
 
-      CNNTstepEncoder
-      FFTstepEncoder
       TstepEncoder
+      FFTstepEncoder
+      CNNTstepEncoder
 

pyproject.toml

@@ -0,0 +1,58 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "amago"
+version = "3.2.0"
+description = "off-policy RL on long sequences"
+readme = { file = "README.md", content-type = "text/markdown" }
+authors = [
+  { name = "Jake Grigsby", email = "grigsby@cs.utexas.edu" }
+]
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Science/Research",
+  "Programming Language :: Python",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+keywords = ["reinforcement-learning", "memory", "transformers", "machine-learning"]
+dependencies = [
+  "gymnasium>=0.26,<=0.29.1",
+  "torch>=2.5,<3",
+  "numpy>=1.21,<3",
+  "gin-config",
+  "wandb",
+  "einops",
+  "tqdm",
+  "gym",
+  "accelerate>=1.0",
+  "termcolor",
+]
+
+[project.optional-dependencies]
+envs = [
+  "metaworld @ git+https://github.com/Farama-Foundation/Metaworld.git@04be337a12305e393c0caf0cbf5ec7755c7c8feb",
+  "cython<3",
+  "popgym",
+  "minigrid",
+  "stable-retro",
+  "dm_env",
+  "dm_alchemy @ git+https://github.com/deepmind/dm_alchemy.git",
+  "gymnax",
+  "matplotlib",
+  "opencv-python",
+  "procgen",
+  "ale_py>=0.10",
+]
+flash = ["ninja", "packaging", "flash-attn"]
+mamba = ["causal-conv1d>=1.1.0", "mamba-ssm"]
+
+[project.urls]
+Repository = "https://github.com/UT-Austin-RPL/amago"
+
+[tool.setuptools.packages.find]
+include = ["amago*"]