Apply upgraded black format

d3rlpy/algos/qlearning/awac.py

@@ -71,6 +71,7 @@ class AWACConfig(LearnableConfig):
             :math:`A^\pi(s_t, a_t)`.
         n_critics (int): Number of Q functions for ensemble.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     actor_optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/qlearning/bc.py

@@ -50,6 +50,7 @@ class BCConfig(LearnableConfig):
             Observation preprocessor.
         action_scaler (d3rlpy.preprocessing.ActionScaler): Action preprocessor.
     """
+
     batch_size: int = 100
     learning_rate: float = 1e-3
     policy_type: str = "deterministic"
@@ -133,6 +134,7 @@ class DiscreteBCConfig(LearnableConfig):
         observation_scaler (d3rlpy.preprocessing.ObservationScaler):
             Observation preprocessor.
     """
+
     batch_size: int = 100
     learning_rate: float = 1e-3
     optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/qlearning/bcq.py

@@ -138,6 +138,7 @@ class BCQConfig(LearnableConfig):
             functions. If this is large, RL training would be more stabilized.
         beta (float): KL reguralization term for Conditional VAE.
     """
+
     actor_learning_rate: float = 1e-3
     critic_learning_rate: float = 1e-3
     imitator_learning_rate: float = 1e-3
@@ -323,6 +324,7 @@ class DiscreteBCQConfig(LearnableConfig):
         share_encoder (bool): Flag to share encoder between Q-function and
             imitation models.
     """
+
     learning_rate: float = 6.25e-5
     optim_factory: OptimizerFactory = make_optimizer_field()
     encoder_factory: EncoderFactory = make_encoder_field()

d3rlpy/algos/qlearning/bear.py

@@ -115,6 +115,7 @@ class BEARConfig(LearnableConfig):
         warmup_steps (int): Number of steps to warmup the policy
             function.
     """
+
     actor_learning_rate: float = 1e-4
     critic_learning_rate: float = 3e-4
     imitator_learning_rate: float = 3e-4

d3rlpy/algos/qlearning/cql.py

@@ -100,6 +100,7 @@ class CQLConfig(LearnableConfig):
             :math:`\log{\sum_a \exp{Q(s, a)}}`.
         soft_q_backup (bool): Flag to use SAC-style backup.
     """
+
     actor_learning_rate: float = 1e-4
     critic_learning_rate: float = 3e-4
     temp_learning_rate: float = 1e-4
@@ -256,6 +257,7 @@ class DiscreteCQLConfig(LearnableConfig):
             network.
         alpha (float): math:`\alpha` value above.
     """
+
     learning_rate: float = 6.25e-5
     optim_factory: OptimizerFactory = make_optimizer_field()
     encoder_factory: EncoderFactory = make_encoder_field()

d3rlpy/algos/qlearning/crr.py

@@ -100,6 +100,7 @@ class CRRConfig(LearnableConfig):
         update_actor_interval (int): Interval to update policy function used
             with ``hard`` target update.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     actor_optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/qlearning/ddpg.py

@@ -70,6 +70,7 @@ class DDPGConfig(LearnableConfig):
         tau (float): Target network synchronization coefficiency.
         n_critics (int): Number of Q functions for ensemble.
     """
+
     batch_size: int = 256
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4

d3rlpy/algos/qlearning/dqn.py

@@ -45,6 +45,7 @@ class DQNConfig(LearnableConfig):
         n_critics (int): Number of Q functions for ensemble.
         target_update_interval (int): Interval to update the target network.
     """
+
     batch_size: int = 32
     learning_rate: float = 6.25e-5
     optim_factory: OptimizerFactory = make_optimizer_field()
@@ -147,6 +148,7 @@ class DoubleDQNConfig(DQNConfig):
         target_update_interval (int): Interval to synchronize the target
             network.
     """
+
     batch_size: int = 32
     learning_rate: float = 6.25e-5
     optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/qlearning/iql.py

@@ -81,6 +81,7 @@ class IQLConfig(LearnableConfig):
             :math:`\beta`.
         max_weight (float): Maximum advantage weight value to clip.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     actor_optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/qlearning/nfq.py

@@ -48,6 +48,7 @@ class NFQConfig(LearnableConfig):
         gamma (float): Discount factor.
         n_critics (int): Number of Q functions for ensemble.
     """
+
     learning_rate: float = 6.25e-5
     optim_factory: OptimizerFactory = make_optimizer_field()
     encoder_factory: EncoderFactory = make_encoder_field()

d3rlpy/algos/qlearning/plas.py

@@ -78,6 +78,7 @@ class PLASConfig(LearnableConfig):
         warmup_steps (int): Number of steps to warmup the VAE.
         beta (float): KL reguralization term for Conditional VAE.
     """
+
     actor_learning_rate: float = 1e-4
     critic_learning_rate: float = 1e-3
     imitator_learning_rate: float = 1e-4
@@ -239,6 +240,7 @@ class PLASWithPerturbationConfig(PLASConfig):
         warmup_steps (int): Number of steps to warmup the VAE.
         beta (float): KL reguralization term for Conditional VAE.
     """
+
     action_flexibility: float = 0.05
 
     def create(self, device: DeviceArg = False) -> "PLASWithPerturbation":

d3rlpy/algos/qlearning/random_policy.py

@@ -31,6 +31,7 @@ class RandomPolicyConfig(LearnableConfig):
         normal_std (float): Standard deviation of the normal distribution. This
             is only used when ``distribution='normal'``.
     """
+
     distribution: str = "uniform"
     normal_std: float = 1.0
 

d3rlpy/algos/qlearning/sac.py

@@ -95,6 +95,7 @@ class SACConfig(LearnableConfig):
         n_critics (int): Number of Q functions for ensemble.
         initial_temperature (float): Initial temperature value.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     temp_learning_rate: float = 3e-4
@@ -243,6 +244,7 @@ class DiscreteSACConfig(LearnableConfig):
         n_critics (int): Number of Q functions for ensemble.
         initial_temperature (float): Initial temperature value.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     temp_learning_rate: float = 3e-4

d3rlpy/algos/qlearning/td3.py

@@ -75,6 +75,7 @@ class TD3Config(LearnableConfig):
         update_actor_interval (int): Interval to update policy function
             described as `delayed policy update` in the paper.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     actor_optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/qlearning/td3_plus_bc.py

@@ -66,6 +66,7 @@ class TD3PlusBCConfig(LearnableConfig):
         update_actor_interval (int): Interval to update policy function
             described as `delayed policy update` in the paper.
     """
+
     actor_learning_rate: float = 3e-4
     critic_learning_rate: float = 3e-4
     actor_optim_factory: OptimizerFactory = make_optimizer_field()

d3rlpy/algos/transformer/action_samplers.py

@@ -48,6 +48,7 @@ class SoftmaxTransformerActionSampler(TransformerActionSampler):
     Args:
         temperature (int): Softmax temperature.
     """
+
     _temperature: float
 
     def __init__(self, temperature: float = 1.0):

d3rlpy/algos/transformer/base.py

@@ -46,7 +46,7 @@ def predict(self, inpt: TorchTransformerInput) -> torch.Tensor:
 
     @abstractmethod
     def inner_predict(self, inpt: TorchTransformerInput) -> torch.Tensor:
-        ...
+        raise NotImplementedError
 
     @train_api
     def update(
@@ -58,7 +58,7 @@ def update(
     def inner_update(
         self, batch: TorchTrajectoryMiniBatch, grad_step: int
     ) -> Dict[str, float]:
-        pass
+        raise NotImplementedError
 
 
 @dataclasses.dataclass()
@@ -101,6 +101,7 @@ class StatefulTransformerWrapper(Generic[TTransformerImpl, TTransformerConfig]):
         target_return (float): Target return.
         action_sampler (d3rlpy.algos.TransformerActionSampler): Action sampler.
     """
+
     _algo: "TransformerAlgoBase[TTransformerImpl, TTransformerConfig]"
     _target_return: float
     _action_sampler: TransformerActionSampler

d3rlpy/base.py

@@ -97,9 +97,9 @@ def unwrap_models_by_ddp(self) -> None:
 class LearnableConfig(DynamicConfig):
     batch_size: int = 256
     gamma: float = 0.99
-    observation_scaler: Optional[
-        ObservationScaler
-    ] = make_observation_scaler_field()
+    observation_scaler: Optional[ObservationScaler] = (
+        make_observation_scaler_field()
+    )
     action_scaler: Optional[ActionScaler] = make_action_scaler_field()
     reward_scaler: Optional[RewardScaler] = make_reward_scaler_field()
 

d3rlpy/dataset/buffers.py

@@ -44,6 +44,7 @@ def __getitem__(self, index: int) -> Tuple[EpisodeBase, int]:
 
 class InfiniteBuffer(BufferProtocol):
     r"""Buffer with unlimited capacity."""
+
     _transitions: List[Tuple[EpisodeBase, int]]
     _episodes: List[EpisodeBase]
 
@@ -78,6 +79,7 @@ class FIFOBuffer(BufferProtocol):
     Args:
         limit (int): buffer capacity.
     """
+
     _transitions: Deque[Tuple[EpisodeBase, int]]
     _episodes: List[EpisodeBase]
     _limit: int

d3rlpy/dataset/components.py

@@ -40,6 +40,7 @@ class Signature:
         dtype: List of numpy data types.
         shape: List of array shapes.
     """
+
     dtype: Sequence[DType]
     shape: Sequence[Sequence[int]]
 
@@ -69,6 +70,7 @@ class Transition:
         terminal: Flag of environment termination.
         interval: Timesteps between ``observation`` and ``next_observation``.
     """
+
     observation: Observation  # (...)
     action: NDArray  # (...)
     reward: Float32NDArray  # (1,)
@@ -130,6 +132,7 @@ class PartialTrajectory:
         masks: Sequence of masks that represent padding.
         length: Sequence length.
     """
+
     observations: ObservationSequence  # (L, ...)
     actions: NDArray  # (L, ...)
     rewards: Float32NDArray  # (L, 1)
@@ -350,6 +353,7 @@ class Episode:
         rewards: Sequence of rewards.
         terminated: Flag of environment termination.
     """
+
     observations: ObservationSequence
     actions: NDArray
     rewards: Float32NDArray
@@ -422,6 +426,7 @@ class DatasetInfo:
             this represents dimension of action vectors. For discrete
             action-space, this represents the number of discrete actions.
     """
+
     observation_signature: Signature
     action_signature: Signature
     reward_signature: Signature

d3rlpy/dataset/episode_generator.py

@@ -32,6 +32,7 @@ class EpisodeGenerator(EpisodeGeneratorProtocol):
         terminals: Sequence of environment terminal flags.
         timeouts: Sequence of timeout flags.
     """
+
     _observations: ObservationSequence
     _actions: NDArray
     _rewards: Float32NDArray

d3rlpy/dataset/mini_batch.py

@@ -30,6 +30,7 @@ class TransitionMiniBatch:
         intervals: Batched timesteps between observations and next
             observations.
     """
+
     observations: Union[Float32NDArray, Sequence[Float32NDArray]]  # (B, ...)
     actions: Float32NDArray  # (B, ...)
     rewards: Float32NDArray  # (B, 1)
@@ -146,6 +147,7 @@ class TrajectoryMiniBatch:
         masks: Batched masks that represent padding.
         length: Length of trajectories.
     """
+
     observations: Union[Float32NDArray, Sequence[Float32NDArray]]  # (B, L, ...)
     actions: Float32NDArray  # (B, L, ...)
     rewards: Float32NDArray  # (B, L, 1)

d3rlpy/dataset/replay_buffer.py

@@ -329,6 +329,7 @@ class ReplayBuffer(ReplayBufferBase):
             for online training. ``cache_size`` needs to be greater than the
             maximum possible episode length.
     """
+
     _buffer: BufferProtocol
     _transition_picker: TransitionPickerProtocol
     _trajectory_slicer: TrajectorySlicerProtocol
@@ -587,6 +588,7 @@ class MixedReplayBuffer(ReplayBufferBase):
         secondary_mix_ratio (float): Ratio to sample mini-batches from the
             secondary replay buffer.
     """
+
     _primary_replay_buffer: ReplayBufferBase
     _secondary_replay_buffer: ReplayBufferBase
     _secondary_mix_ratio: float

d3rlpy/dataset/trajectory_slicers.py

@@ -116,6 +116,7 @@ class FrameStackTrajectorySlicer(TrajectorySlicerProtocol):
     Args:
         n_frames: Number of frames to stack.
     """
+
     _n_frames: int
 
     def __init__(self, n_frames: int):

d3rlpy/dataset/transition_pickers.py

@@ -103,6 +103,7 @@ class FrameStackTransitionPicker(TransitionPickerProtocol):
         n_frames (int): Number of frames to stack.
         gamma (float): Discount factor to compute return-to-go.
     """
+
     _n_frames: int
     _gamma: float
 
@@ -152,6 +153,7 @@ class MultiStepTransitionPicker(TransitionPickerProtocol):
             ``net_observation``.
         gamma: Discount factor to compute a multi-step return.
     """
+
     _n_steps: int
     _gamma: float