rlgym-learn-algos 0.2.1__cp310-cp310-musllinux_1_2_aarch64.whl → 0.2.3__cp310-cp310-musllinux_1_2_aarch64.whl

rlgym_learn_algos/ppo/experience_buffer.py

@@ -7,6 +7,7 @@ import numpy as np
 import torch
 from pydantic import BaseModel, Field, model_validator
 from rlgym.api import ActionType, AgentID, ObsType, RewardType
+
 from rlgym_learn_algos.util.torch_functions import get_device
 from rlgym_learn_algos.util.torch_pydantic import PydanticTorchDevice
 
@@ -24,6 +25,7 @@ EXPERIENCE_BUFFER_FILE = "experience_buffer.pkl"
 class ExperienceBufferConfigModel(BaseModel, extra="forbid"):
     max_size: int = 100000
     device: PydanticTorchDevice = "auto"
+    save_experience_buffer_in_checkpoint: bool = True
     trajectory_processor_config: Dict[str, Any] = Field(default_factory=dict)
 
     @model_validator(mode="before")
@@ -40,8 +42,9 @@ class ExperienceBufferConfigModel(BaseModel, extra="forbid"):
                     data["trajectory_processor_config"] = data[
                         "trajectory_processor_config"
                     ].model_dump()
-            if "device" not in data or data["device"] == "auto":
-                data["device"] = get_device("auto")
+            if "device" not in data:
+                data["device"] = "auto"
+            data["device"] = get_device(data["device"])
         return data
 
 
@@ -165,21 +168,22 @@ class ExperienceBuffer(
 
     def save_checkpoint(self, folder_path):
         os.makedirs(folder_path, exist_ok=True)
-        with open(
-            os.path.join(folder_path, EXPERIENCE_BUFFER_FILE),
-            "wb",
-        ) as f:
-            pickle.dump(
-                {
-                    "agent_ids": self.agent_ids,
-                    "observations": self.observations,
-                    "actions": self.actions,
-                    "log_probs": self.log_probs,
-                    "values": self.values,
-                    "advantages": self.advantages,
-                },
-                f,
-            )
+        if self.config.experience_buffer_config.save_experience_buffer_in_checkpoint:
+            with open(
+                os.path.join(folder_path, EXPERIENCE_BUFFER_FILE),
+                "wb",
+            ) as f:
+                pickle.dump(
+                    {
+                        "agent_ids": self.agent_ids,
+                        "observations": self.observations,
+                        "actions": self.actions,
+                        "log_probs": self.log_probs,
+                        "values": self.values,
+                        "advantages": self.advantages,
+                    },
+                    f,
+                )
         self.trajectory_processor.save_checkpoint(folder_path)
 
     # TODO: update docs

rlgym_learn_algos/ppo/ppo_agent_controller.py

@@ -24,6 +24,8 @@ from rlgym.api import (
 )
 from rlgym_learn import EnvActionResponse, EnvActionResponseType, Timestep
 from rlgym_learn.api.agent_controller import AgentController
+from torch import device as _device
+
 from rlgym_learn_algos.logging import (
     DerivedMetricsLoggerConfig,
     MetricsLogger,
@@ -34,7 +36,6 @@ from rlgym_learn_algos.logging import (
 )
 from rlgym_learn_algos.stateful_functions import ObsStandardizer
 from rlgym_learn_algos.util.torch_functions import get_device
-from torch import device as _device
 
 from .actor import Actor
 from .critic import Critic
@@ -57,8 +58,8 @@ EXPERIENCE_BUFFER_FOLDER = "experience_buffer"
 PPO_LEARNER_FOLDER = "ppo_learner"
 METRICS_LOGGER_FOLDER = "metrics_logger"
 PPO_AGENT_FILE = "ppo_agent.json"
+ITERATION_TRAJECTORIES_FILE = "current_trajectories.pkl"  # this should be renamed, but it would be a breaking change so I'm leaving it until I happen to make one of those and remember to update this at the same time
 ITERATION_SHARED_INFOS_FILE = "iteration_shared_infos.pkl"
-CURRENT_TRAJECTORIES_FILE = "current_trajectories.pkl"
 
 
 class PPOAgentControllerConfigModel(BaseModel, extra="forbid"):
@@ -68,6 +69,7 @@ class PPOAgentControllerConfigModel(BaseModel, extra="forbid"):
     checkpoint_load_folder: Optional[str] = None
     n_checkpoints_to_keep: int = 5
     random_seed: int = 123
+    save_mid_iteration_data_in_checkpoint = True
     learner_config: PPOLearnerConfigModel = Field(default_factory=PPOLearnerConfigModel)
     experience_buffer_config: ExperienceBufferConfigModel = Field(
         default_factory=ExperienceBufferConfigModel
@@ -166,7 +168,7 @@ class PPOAgentController(
             str,
             EnvTrajectories[AgentID, ActionType, ObsType, RewardType],
         ] = {}
-        self.current_trajectories: List[
+        self.iteration_trajectories: List[
             Trajectory[AgentID, ActionType, ObsType, RewardType]
         ] = []
         self.iteration_shared_infos: List[Dict[str, Any]] = []
@@ -309,19 +311,18 @@ class PPOAgentController(
             with open(
                 os.path.join(
                     self.config.agent_controller_config.checkpoint_load_folder,
-                    CURRENT_TRAJECTORIES_FILE,
+                    ITERATION_TRAJECTORIES_FILE,
                 ),
                 "rb",
             ) as f:
-                current_trajectories: Dict[
-                    int,
-                    EnvTrajectories[AgentID, ActionType, ObsType, RewardType],
+                iteration_trajectories: List[
+                    Trajectory[AgentID, ObsType, ActionType, RewardType]
                 ] = pickle.load(f)
         except FileNotFoundError:
             print(
-                f"{self.config.agent_controller_name}: Tried to load current trajectories from checkpoint using the file at location {str(os.path.join(self.config.agent_controller_config.checkpoint_load_folder, CURRENT_TRAJECTORIES_FILE))}, but there is no such file! Current trajectories will be initialized as an empty dict instead."
+                f"{self.config.agent_controller_name}: Tried to load current trajectories from checkpoint using the file at location {str(os.path.join(self.config.agent_controller_config.checkpoint_load_folder, ITERATION_TRAJECTORIES_FILE))}, but there is no such file! Current trajectories will be initialized as an empty list instead."
             )
-            current_trajectories = {}
+            iteration_trajectories = []
         try:
             with open(
                 os.path.join(
@@ -335,7 +336,7 @@ class PPOAgentController(
             print(
                 f"{self.config.agent_controller_name}: Tried to load iteration shared info data from checkpoint using the file at location {str(os.path.join(self.config.agent_controller_config.checkpoint_load_folder, ITERATION_SHARED_INFOS_FILE))}, but there is no such file! Iteration shared info data will be initialized as an empty list instead."
             )
-            current_trajectories = {}
+            iteration_shared_infos = []
         try:
             with open(
                 os.path.join(
@@ -357,7 +358,7 @@ class PPOAgentController(
                 "timestep_collection_start_time": time.perf_counter(),
             }
 
-        self.current_trajectories = current_trajectories
+        self.iteration_trajectories = iteration_trajectories
         self.iteration_shared_infos = iteration_shared_infos
         self.cur_iteration = state["cur_iteration"]
         self.iteration_timesteps = state["iteration_timesteps"]
@@ -379,20 +380,22 @@ class PPOAgentController(
         self.experience_buffer.save_checkpoint(
             os.path.join(checkpoint_save_folder, EXPERIENCE_BUFFER_FOLDER)
         )
-        self.metrics_logger.save_checkpoint(
-            os.path.join(checkpoint_save_folder, METRICS_LOGGER_FOLDER)
-        )
+        if self.metrics_logger is not None:
+            self.metrics_logger.save_checkpoint(
+                os.path.join(checkpoint_save_folder, METRICS_LOGGER_FOLDER)
+            )
 
-        with open(
-            os.path.join(checkpoint_save_folder, CURRENT_TRAJECTORIES_FILE),
-            "wb",
-        ) as f:
-            pickle.dump(self.current_trajectories, f)
-        with open(
-            os.path.join(checkpoint_save_folder, ITERATION_SHARED_INFOS_FILE),
-            "wb",
-        ) as f:
-            pickle.dump(self.iteration_shared_infos, f)
+        if self.config.agent_controller_config.save_mid_iteration_data_in_checkpoint:
+            with open(
+                os.path.join(checkpoint_save_folder, ITERATION_TRAJECTORIES_FILE),
+                "wb",
+            ) as f:
+                pickle.dump(self.iteration_trajectories, f)
+            with open(
+                os.path.join(checkpoint_save_folder, ITERATION_SHARED_INFOS_FILE),
+                "wb",
+            ) as f:
+                pickle.dump(self.iteration_shared_infos, f)
         with open(os.path.join(checkpoint_save_folder, PPO_AGENT_FILE), "wt") as f:
             state = {
                 "cur_iteration": self.cur_iteration,
@@ -403,6 +406,7 @@ class PPOAgentController(
             }
             json.dump(state, f, indent=4)
 
+        # TODO: does this actually work? I'm not sure the file structure I'm using actually works with this assumption
         # Prune old checkpoints
         existing_checkpoints = [
             int(arg) for arg in os.listdir(self.checkpoints_save_folder)
@@ -506,7 +510,7 @@ class PPOAgentController(
             elif enum_type == EnvActionResponseType.RESET:
                 env_trajectories = self.current_env_trajectories.pop(env_id)
                 env_trajectories.finalize()
-                self.current_trajectories += env_trajectories.get_trajectories()
+                self.iteration_trajectories += env_trajectories.get_trajectories()
             elif enum_type == EnvActionResponseType.SET_STATE:
                 # Can get the desired_state using env_action.desired_state and the prev_timestep_id_dict using env_action.prev_timestep_id_dict, but I'll leave that to you
                 raise NotImplementedError
@@ -517,10 +521,10 @@ class PPOAgentController(
         env_trajectories_list = list(self.current_env_trajectories.values())
         for env_trajectories in env_trajectories_list:
             env_trajectories.finalize()
-            self.current_trajectories += env_trajectories.get_trajectories()
+            self.iteration_trajectories += env_trajectories.get_trajectories()
         self._update_value_predictions()
         trajectory_processor_data = self.experience_buffer.submit_experience(
-            self.current_trajectories
+            self.iteration_trajectories
         )
         ppo_data = self.learner.learn(self.experience_buffer)
 
@@ -540,9 +544,9 @@ class PPOAgentController(
             self.metrics_logger.collect_env_metrics(self.iteration_shared_infos)
             self.metrics_logger.report_metrics()
 
-        self.iteration_shared_infos = []
+        self.iteration_trajectories.clear()
+        self.iteration_shared_infos.clear()
         self.current_env_trajectories.clear()
-        self.current_trajectories.clear()
         self.ts_since_last_save += self.iteration_timesteps
         self.iteration_timesteps = 0
         self.iteration_start_time = cur_time
@@ -551,14 +555,14 @@ class PPOAgentController(
     @torch.no_grad()
     def _update_value_predictions(self):
         """
-        Function to update the value predictions inside the Trajectory instances of self.current_trajectories
+        Function to update the value predictions inside the Trajectory instances of self.iteration_trajectories
         """
         traj_timestep_idx_ranges: List[Tuple[int, int]] = []
         start = 0
         stop = 0
         critic_agent_id_input: List[AgentID] = []
         critic_obs_input: List[ObsType] = []
-        for trajectory in self.current_trajectories:
+        for trajectory in self.iteration_trajectories:
             obs_list = trajectory.obs_list + [trajectory.final_obs]
             traj_len = len(obs_list)
             agent_id_list = [trajectory.agent_id] * traj_len
@@ -575,7 +579,7 @@ class PPOAgentController(
         )
         torch.cuda.empty_cache()
         for idx, (start, stop) in enumerate(traj_timestep_idx_ranges):
-            self.current_trajectories[idx].val_preds = val_preds[start : stop - 1]
-            self.current_trajectories[idx].final_val_pred = val_preds[stop - 1]
+            self.iteration_trajectories[idx].val_preds = val_preds[start : stop - 1]
+            self.iteration_trajectories[idx].final_val_pred = val_preds[stop - 1]
         if self.config.agent_controller_config.learner_config.device.type != "cpu":
             torch.cuda.current_stream().synchronize()

rlgym_learn_algos/ppo/ppo_learner.py

@@ -39,17 +39,25 @@ class PPOLearnerConfigModel(BaseModel, extra="forbid"):
     clip_range: float = 0.2
     actor_lr: float = 3e-4
     critic_lr: float = 3e-4
+    advantage_normalization: bool = True
     device: PydanticTorchDevice = "auto"
+    cudnn_benchmark_mode: bool = True
 
     @model_validator(mode="before")
     @classmethod
     def set_device(cls, data):
-        if isinstance(data, dict) and (
-            "device" not in data or data["device"] == "auto"
-        ):
-            data["device"] = get_device("auto")
+        if isinstance(data, dict):
+            if "device" not in data:
+                data["device"] = "auto"
+            data["device"] = get_device(data["device"])
         return data
 
+    @model_validator(mode="after")
+    def validate_cudnn_benchmark(self):
+        if self.device.type != "cuda":
+            self.cudnn_benchmark_mode = False
+        return self
+
 
 @dataclass
 class DerivedPPOLearnerConfig:
@@ -107,6 +115,12 @@ class PPOLearner(
     def load(self, config: DerivedPPOLearnerConfig):
         self.config = config
 
+        if (
+            config.learner_config.cudnn_benchmark_mode
+            and config.learner_config.device.type == "cuda"
+        ):
+            torch.backends.cudnn.benchmark = True
+
         self.actor = self.actor_factory(
             config.obs_space, config.action_space, config.learner_config.device
         )
@@ -292,6 +306,10 @@ class PPOLearner(
                     advantages = batch_advantages[start:stop].to(
                         self.config.learner_config.device
                     )
+                    if self.config.learner_config.advantage_normalization:
+                        advantages = (advantages - torch.mean(advantages)) / (
+                            torch.std(advantages) + 1e-8
+                        )
                     old_probs = batch_old_probs[start:stop].to(
                         self.config.learner_config.device
                     )

rlgym_learn_algos/util/torch_functions.py

@@ -13,7 +13,6 @@ import torch.nn as nn
 
 def get_device(device: str):
     if device in ["auto", "gpu"] and torch.cuda.is_available():
-        torch.backends.cudnn.benchmark = True
         return "cuda:0"
     elif device == "auto" and not torch.cuda.is_available():
         return "cpu"

rlgym_learn_algos/util/torch_pydantic.py

@@ -42,7 +42,7 @@ device_str_regex = (
             "privateuseone",
         ]
     )
-    + ")(:\d+)"
+    + ")(:\d+)?"
 )
 
 

{rlgym_learn_algos-0.2.1.dist-info → rlgym_learn_algos-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rlgym-learn-algos
-Version: 0.2.1
+Version: 0.2.3
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Requires-Dist: pydantic>=2.8.2

{rlgym_learn_algos-0.2.1.dist-info → rlgym_learn_algos-0.2.3.dist-info}/RECORD

@@ -1,6 +1,6 @@
-rlgym_learn_algos-0.2.1.dist-info/METADATA,sha256=u9tM8KxVpm3nQ86a01tlDy5csGAeTmCMDZdccXgrUrg,2403
-rlgym_learn_algos-0.2.1.dist-info/WHEEL,sha256=uRJ-CvahRmYoh93Vk8bjAlQe2Efq6pkWnXHZVi-TE-E,108
-rlgym_learn_algos-0.2.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+rlgym_learn_algos-0.2.3.dist-info/METADATA,sha256=tJnTSp4XVAuNddScvlgboTLPico7e8xFGA6XAdVfLVc,2403
+rlgym_learn_algos-0.2.3.dist-info/WHEEL,sha256=ST2opB6BOOjt6oQEfC3w015mY7a2mgvvpUSxQ5LHNYE,108
+rlgym_learn_algos-0.2.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 rlgym_learn_algos.libs/libgcc_s-e52197c3.so.1,sha256=vkPW1Auw6CH9Bjk7frmX3hry_1H9c0tRI0Ncyg71WUI,724137
 rlgym_learn_algos/__init__.py,sha256=dZeTgNro6qG1Hu0l0UBhgHOYiyeCwPWndC84dJAp__U,203
 rlgym_learn_algos/conversion/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,18 +16,18 @@ rlgym_learn_algos/ppo/continuous_actor.py,sha256=A4FQ0lKqlB47AeSrDdrPXMvNKXhl5to
 rlgym_learn_algos/ppo/critic.py,sha256=XPleWDO8uM25zlzptWDvZQpUpKvib5kRs9JpmWTVPuY,669
 rlgym_learn_algos/ppo/discrete_actor.py,sha256=TZC7b7ss16giobPC1oz-maOSDX-SrNBUzS1wIV2Rzgw,2547
 rlgym_learn_algos/ppo/env_trajectories.py,sha256=PaO6dmpNkQ3yDLaHIRc0ipn45t5zAjE5U1D_N-LQtgY,3684
-rlgym_learn_algos/ppo/experience_buffer.py,sha256=f_baTo18JKBQjTQoO0FJoCOWlZFV8vJgMXqevg-TSi8,10552
+rlgym_learn_algos/ppo/experience_buffer.py,sha256=4wWSfq1tobXv7lmBbkM5sbTVuPJxrdAbxi5rNHc04g4,10769
 rlgym_learn_algos/ppo/experience_buffer_numpy.py,sha256=6AOGQjDn_dHLS9bmxJW_cGEjBUbe8u5VWS0LVlpIdmY,4617
 rlgym_learn_algos/ppo/gae_trajectory_processor.py,sha256=gv5kxvvPnK7SyQIAq6MbOFILIMdPlzoLZwM8TRmtNWw,5302
 rlgym_learn_algos/ppo/gae_trajectory_processor_pure_python.py,sha256=cq7qbK0mcLDXRzA6-pKW0OC50X52XhT5himcOTD6Ei4,6657
 rlgym_learn_algos/ppo/multi_discrete_actor.py,sha256=Mik0X79dUy2ZRIMol4RMTZE9qzsOk6f_6bDaOl5ghxs,3039
-rlgym_learn_algos/ppo/ppo_agent_controller.py,sha256=iVmCvN7H1IcKX7VrZnMParQ148EnPqs9yZ9CCgtdsq8,24524
-rlgym_learn_algos/ppo/ppo_learner.py,sha256=oRDBQszbgFvoVmGjRqpklQWTOI2NazOLzhDU_c9-SAU,14827
+rlgym_learn_algos/ppo/ppo_agent_controller.py,sha256=9A-skvF8Bu9tCNBuatYQ5vG_LaQEEkHZ9lhmMdrVp80,25045
+rlgym_learn_algos/ppo/ppo_learner.py,sha256=utEWkikXCpC6Xc1D3asohO0HsIaq3tLyoTlb7fXLOw4,15522
 rlgym_learn_algos/ppo/ppo_metrics_logger.py,sha256=iUyUc2GPwDIIjZeJPZWxoeRrzUWV_qLOac0vApQBkp0,2803
 rlgym_learn_algos/ppo/trajectory.py,sha256=_xyS9ueU6iVvqMUpFr-kb42wEHHZy4zCse7_r660n5E,690
 rlgym_learn_algos/ppo/trajectory_processor.py,sha256=3XRsXXexHWp6UV5nAeBLYvWqvQ9EbNHSN3Yooi4cezo,2031
 rlgym_learn_algos/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rlgym_learn_algos/rlgym_learn_algos.cpython-310-aarch64-linux-gnu.so,sha256=jV1NFWiG87tgaIASpAn-sYP1kOU9X2Kx9reQPZVTEGA,854865
+rlgym_learn_algos/rlgym_learn_algos.cpython-310-aarch64-linux-gnu.so,sha256=ArAiwVGgVpDWQoRza7ndfAx6SFs9qFSPpUM2HGuovls,854865
 rlgym_learn_algos/rlgym_learn_algos.pyi,sha256=B9Kt9uK8xCqASRxWvzLdV501TSCMO4vTNqvZ0MhOHyo,1164
 rlgym_learn_algos/stateful_functions/__init__.py,sha256=OAVy6cQIS85Utyp18jjHgdmascX_8nkwk3A0OpFJxT4,230
 rlgym_learn_algos/stateful_functions/batch_reward_type_numpy_converter.py,sha256=a3q2l5SIgDI36ImF_kYoa684pghnFnlV2vGYvV2zcV0,743
@@ -35,6 +35,6 @@ rlgym_learn_algos/stateful_functions/numpy_obs_standardizer.py,sha256=Xa_fuJCSGu
 rlgym_learn_algos/stateful_functions/obs_standardizer.py,sha256=m2nw1JUg2MKYthn6tWrv2HYIuQ-GfUm48RX9s99jXF4,589
 rlgym_learn_algos/util/__init__.py,sha256=hq7M00Q7zAfyQmIGmXOif0vI40aj_FQ5SqI5dnuGvb0,46
 rlgym_learn_algos/util/running_stats.py,sha256=KtzdKKT75-5ZC58JRqaDXk6sBqa3ZSjQQZrRajAw3Yk,4339
-rlgym_learn_algos/util/torch_functions.py,sha256=6esZL8FeVwWWQWvwLEhkh_B0WqfIWzIWkCDbuhmRlnk,3366
-rlgym_learn_algos/util/torch_pydantic.py,sha256=zXllJoV8HgqJxguPKJ4Y3DIWEwDeJlBW9CIps-yxM44,3412
-rlgym_learn_algos-0.2.1.dist-info/RECORD,,
+rlgym_learn_algos/util/torch_functions.py,sha256=ImgDw4I3ZixGDi17YRkW6UbaiaQTbvOCUCS7N0QVSsU,3320
+rlgym_learn_algos/util/torch_pydantic.py,sha256=khPGA6kWh4_WHoploDkl_SCIGX8SkKkFT40RE06PImc,3413
+rlgym_learn_algos-0.2.3.dist-info/RECORD,,

{rlgym_learn_algos-0.2.1.dist-info → rlgym_learn_algos-0.2.3.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.8.6)
+Generator: maturin (1.9.0)
 Root-Is-Purelib: false
 Tag: cp310-cp310-musllinux_1_2_aarch64