fix: formatting

Makefile

@@ -1,4 +1,4 @@
-.PHONY: clean clean-build clean-pyc clean-test coverage dist docs help install check formatbump-version release
+.PHONY: clean clean-build clean-pyc clean-test coverage dist docs help install check bump-version release format
 .DEFAULT_GOAL := help
 
 define BROWSER_PYSCRIPT

README.md

@@ -7,7 +7,7 @@
  </a>
 </p>
 
-[![PyPI Version](https://img.shields.io/pypi/v/arlbench.svg)](https://pypi.python.org/pypi/arlbench)
+<!--- [![PyPI Version](https://img.shields.io/pypi/v/arlbench.svg)](https://pypi.python.org/pypi/arlbench) -->
 [![Test](https://github.com/automl/arlbench/actions/workflows/pytest.yaml/badge.svg)](https://github.com/automl/arlbench/actions/workflows/pytest.yaml)
 [![Doc Status](https://github.com/automl/arlbench/actions/workflows/docs.yaml/badge.svg)](https://github.com/automl/arlbench/actions/workflows/docs.yaml)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

arlbench/autorl/autorl_env.py

@@ -9,15 +9,8 @@
 import numpy as np
 import pandas as pd
 
-from arlbench.core.algorithms import (
- DQN,
- PPO,
- SAC,
- Algorithm,
- AlgorithmState,
- TrainResult,
- TrainReturnT,
-)
+from arlbench.core.algorithms import (DQN, PPO, SAC, Algorithm, AlgorithmState,
+ TrainResult, TrainReturnT)
 from arlbench.core.environments import make_env
 from arlbench.utils import config_space_to_gymnasium_space
 

arlbench/autorl/checkpointing.py

@@ -10,9 +10,8 @@
 import jax.numpy as jnp
 import numpy as np
 import orbax.checkpoint as ocp
-from flashbax.buffers.prioritised_trajectory_buffer import (
- PrioritisedTrajectoryBufferState,
-)
+from flashbax.buffers.prioritised_trajectory_buffer import \
+ PrioritisedTrajectoryBufferState
 from flashbax.buffers.sum_tree import SumTreeState
 from flashbax.vault import Vault
 from flax.core.frozen_dict import FrozenDict

arlbench/core/algorithms/__init__.py

@@ -1,35 +1,16 @@
 from collections.abc import Callable
 from typing import Optional, Union
 
-from flashbax.buffers.prioritised_trajectory_buffer import (
- PrioritisedTrajectoryBufferState,
-)
+from flashbax.buffers.prioritised_trajectory_buffer import \
+ PrioritisedTrajectoryBufferState
 
 from .algorithm import Algorithm
-from .dqn import (
- DQN,
- DQNMetrics,
- DQNRunnerState,
- DQNState,
- DQNTrainingResult,
- DQNTrainReturnT,
-)
-from .ppo import (
- PPO,
- PPOMetrics,
- PPORunnerState,
- PPOState,
- PPOTrainingResult,
- PPOTrainReturnT,
-)
-from .sac import (
- SAC,
- SACMetrics,
- SACRunnerState,
- SACState,
- SACTrainingResult,
- SACTrainReturnT,
-)
+from .dqn import (DQN, DQNMetrics, DQNRunnerState, DQNState, DQNTrainingResult,
+ DQNTrainReturnT)
+from .ppo import (PPO, PPOMetrics, PPORunnerState, PPOState, PPOTrainingResult,
+ PPOTrainReturnT)
+from .sac import (SAC, SACMetrics, SACRunnerState, SACState, SACTrainingResult,
+ SACTrainReturnT)
 
 TrainResult = Union[DQNTrainingResult, PPOTrainingResult, SACTrainingResult]
 TrainMetrics = Union[DQNMetrics, PPOMetrics, SACMetrics]

arlbench/core/algorithms/buffers.py

@@ -8,12 +8,9 @@
 from flashbax import utils
 from flashbax.buffers import sum_tree
 from flashbax.buffers.prioritised_trajectory_buffer import (
- Experience,
- PrioritisedTrajectoryBufferSample,
- PrioritisedTrajectoryBufferState,
- _get_sample_trajectories,
- get_invalid_indices,
-)
+ Experience, PrioritisedTrajectoryBufferSample,
+ PrioritisedTrajectoryBufferState, _get_sample_trajectories,
+ get_invalid_indices)
 from flashbax.buffers.trajectory_buffer import calculate_uniform_item_indices
 
 if TYPE_CHECKING:

arlbench/core/algorithms/dqn/__init__.py

@@ -1,11 +1,5 @@
-from .dqn import (
- DQN,
- DQNMetrics,
- DQNRunnerState,
- DQNState,
- DQNTrainingResult,
- DQNTrainReturnT,
-)
+from .dqn import (DQN, DQNMetrics, DQNRunnerState, DQNState, DQNTrainingResult,
+ DQNTrainReturnT)
 
 __all__ = [
  "DQN",

arlbench/core/algorithms/dqn/dqn.py

@@ -10,31 +10,23 @@
 import jax.numpy as jnp
 import numpy as np
 import optax
-from ConfigSpace import (
- Categorical,
- Configuration,
- ConfigurationSpace,
- EqualsCondition,
- Float,
- Integer,
-)
+from ConfigSpace import (Categorical, Configuration, ConfigurationSpace,
+ EqualsCondition, Float, Integer)
 from flax.training.train_state import TrainState
 
 from arlbench.core import running_statistics
 from arlbench.core.algorithms.algorithm import Algorithm
 from arlbench.core.algorithms.buffers import uniform_sample
 from arlbench.core.algorithms.common import TimeStep
-from arlbench.core.algorithms.prioritised_item_buffer import (
- make_prioritised_item_buffer,
-)
+from arlbench.core.algorithms.prioritised_item_buffer import \
+ make_prioritised_item_buffer
 
 from .models import CNNQ, MLPQ
 
 if TYPE_CHECKING:
  import chex
- from flashbax.buffers.prioritised_trajectory_buffer import (
- PrioritisedTrajectoryBufferState,
- )
+ from flashbax.buffers.prioritised_trajectory_buffer import \
+ PrioritisedTrajectoryBufferState
  from flax.core.frozen_dict import FrozenDict
 
  from arlbench.core.environments import Environment

arlbench/core/algorithms/ppo/__init__.py

@@ -1,11 +1,5 @@
-from .ppo import (
- PPO,
- PPOMetrics,
- PPORunnerState,
- PPOState,
- PPOTrainingResult,
- PPOTrainReturnT,
-)
+from .ppo import (PPO, PPOMetrics, PPORunnerState, PPOState, PPOTrainingResult,
+ PPOTrainReturnT)
 
 __all__ = [
  "PPO",

arlbench/core/algorithms/ppo/ppo.py

@@ -10,7 +10,8 @@
 import jax.numpy as jnp
 import numpy as np
 import optax
-from ConfigSpace import Categorical, Configuration, ConfigurationSpace, Float, Integer
+from ConfigSpace import (Categorical, Configuration, ConfigurationSpace, Float,
+ Integer)
 from flax.training.train_state import TrainState
 
 from arlbench.core import running_statistics

arlbench/core/algorithms/prioritised_item_buffer.py

@@ -20,12 +20,9 @@
 from flashbax.buffers.item_buffer import validate_item_buffer_args
 from flashbax.buffers.prioritised_flat_buffer import validate_priority_exponent
 from flashbax.buffers.prioritised_trajectory_buffer import (
- PrioritisedTrajectoryBuffer,
- PrioritisedTrajectoryBufferSample,
- PrioritisedTrajectoryBufferState,
- make_prioritised_trajectory_buffer,
- validate_device,
-)
+ PrioritisedTrajectoryBuffer, PrioritisedTrajectoryBufferSample,
+ PrioritisedTrajectoryBufferState, make_prioritised_trajectory_buffer,
+ validate_device)
 from flashbax.utils import add_dim_to_args
 
 if TYPE_CHECKING:

arlbench/core/algorithms/sac/__init__.py

@@ -1,11 +1,5 @@
-from .sac import (
- SAC,
- SACMetrics,
- SACRunnerState,
- SACState,
- SACTrainingResult,
- SACTrainReturnT,
-)
+from .sac import (SAC, SACMetrics, SACRunnerState, SACState, SACTrainingResult,
+ SACTrainReturnT)
 
 __all__ = [
  "SAC",

arlbench/core/algorithms/sac/sac.py

@@ -10,38 +10,24 @@
 import jax.numpy as jnp
 import numpy as np
 import optax
-from ConfigSpace import (
- Categorical,
- Configuration,
- ConfigurationSpace,
- EqualsCondition,
- Float,
- Integer,
-)
+from ConfigSpace import (Categorical, Configuration, ConfigurationSpace,
+ EqualsCondition, Float, Integer)
 from flax.training.train_state import TrainState
 
 from arlbench.core import running_statistics
 from arlbench.core.algorithms.algorithm import Algorithm
 from arlbench.core.algorithms.buffers import uniform_sample
 from arlbench.core.algorithms.common import TimeStep
-from arlbench.core.algorithms.prioritised_item_buffer import (
- make_prioritised_item_buffer,
-)
-
-from .models import (
- AlphaCoef,
- SACCNNActor,
- SACCNNCritic,
- SACMLPActor,
- SACMLPCritic,
- SACVectorCritic,
-)
+from arlbench.core.algorithms.prioritised_item_buffer import \
+ make_prioritised_item_buffer
+
+from .models import (AlphaCoef, SACCNNActor, SACCNNCritic, SACMLPActor,
+ SACMLPCritic, SACVectorCritic)
 
 if TYPE_CHECKING:
  import chex
- from flashbax.buffers.prioritised_trajectory_buffer import (
- PrioritisedTrajectoryBufferState,
- )
+ from flashbax.buffers.prioritised_trajectory_buffer import \
+ PrioritisedTrajectoryBufferState
  from flax.core.frozen_dict import FrozenDict
 
  from arlbench.core.environments import Environment

arlbench/utils/__init__.py

@@ -1,11 +1,6 @@
-from .common import (
- config_space_to_gymnasium_space,
- config_space_to_yaml,
- gymnasium_space_to_gymnax_space,
- recursive_concat,
- save_defaults_to_yaml,
- tuple_concat,
-)
+from .common import (config_space_to_gymnasium_space, config_space_to_yaml,
+ gymnasium_space_to_gymnax_space, recursive_concat,
+ save_defaults_to_yaml, tuple_concat)
 from .handle_termination import HandleTermination
 
 __all__ = [

tests/autorl/test_autorl_env.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import pytest
+
 from arlbench import AutoRLEnv
 from arlbench.core.algorithms import DQN
 

tests/autorl/test_checkpointing.py

@@ -2,6 +2,7 @@
 
 import jax
 import numpy as np
+
 from arlbench.autorl import AutoRLEnv
 from arlbench.autorl.checkpointing import Checkpointer
 from arlbench.core.algorithms import DQN

tests/autorl/test_objectives.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import numpy as np
+
 from arlbench import AutoRLEnv
 from arlbench.core.algorithms import DQN
 

tests/core/algorithms/test_dqn.py

@@ -4,6 +4,7 @@
 import warnings
 
 import jax
+
 from arlbench.core.algorithms import DQN
 from arlbench.core.environments import make_env
 

tests/core/algorithms/test_ppo.py

@@ -4,6 +4,7 @@
 import warnings
 
 import jax
+
 from arlbench.core.algorithms import PPO
 from arlbench.core.environments import make_env
 

tests/core/algorithms/test_sac.py

@@ -4,6 +4,7 @@
 import warnings
 
 import jax
+
 from arlbench.core.algorithms import SAC
 from arlbench.core.environments import make_env
 

tests/core/environments/test_atari.py

@@ -4,6 +4,7 @@
 
 import jax
 import numpy as np
+
 from arlbench.core.algorithms import DQN, PPO
 from arlbench.core.environments import make_env
 

tests/core/environments/test_gymnasium.py

@@ -4,6 +4,7 @@
 
 import jax
 import numpy as np
+
 from arlbench.core.algorithms import DQN, PPO, SAC
 from arlbench.core.environments import make_env
 

tests/utils/test_utils.py

@@ -2,9 +2,10 @@
 
 import gymnasium
 import numpy as np
-from arlbench.utils import config_space_to_gymnasium_space
 from ConfigSpace import Categorical, ConfigurationSpace, Float, Integer
 
+from arlbench.utils import config_space_to_gymnasium_space
+
 
 def test_config_space_to_gymnasium_space():
  config_space = ConfigurationSpace(