1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115 | """Abstract base class for probability distributions."""
from abc import ABC, abstractmethod
from typing import TypeVar
import torch
from torch import nn
SelfDistribution = TypeVar("SelfDistribution", bound="Distribution")
class Distribution(ABC):
"""Abstract base class for distributions."""
def __init__(self) -> None:
super().__init__()
self.distribution = None
@abstractmethod
def proba_distribution_net(self, *args, **kwargs) -> nn.Module | tuple[nn.Module, nn.Parameter]:
"""Create the layers and parameters that represent the distribution.
Subclasses must define this, but the arguments and return type vary between
concrete classes.
"""
@abstractmethod
def proba_distribution(self: SelfDistribution, *args, **kwargs) -> SelfDistribution:
"""Set parameters of the distribution.
Returns:
self
"""
@abstractmethod
def log_prob(self, x: torch.Tensor) -> torch.Tensor:
"""Returns the log likelihood.
Args:
x: the taken action
Returns:
The log likelihood of the distribution
"""
@abstractmethod
def entropy(self) -> torch.Tensor | None:
"""Returns Shannon's entropy of the probability.
Returns:
the entropy, or None if no analytical form is known
"""
@abstractmethod
def sample(self) -> torch.Tensor:
"""Returns a sample from the probability distribution.
Returns:
the stochastic action
"""
@abstractmethod
def mode(self) -> torch.Tensor:
"""Returns the most likely action (deterministic output).
from the probability distribution.
Returns:
the stochastic action
"""
def get_actions(self, deterministic: bool = False) -> torch.Tensor:
"""Return actions according to the probability distribution.
Args:
deterministic: Whether to use deterministic actions.
"""
if deterministic:
return self.mode()
return self.sample()
@abstractmethod
def actions_from_params(self, *args, **kwargs) -> torch.Tensor:
"""Returns samples from the probability distribution.
given its parameters.
Returns:
actions
"""
@abstractmethod
def log_prob_from_params(self, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
"""Returns samples and the associated log probabilities.
from the probability distribution given its parameters.
Returns:
actions and log prob
"""
def sum_independent_dims(tensor: torch.Tensor) -> torch.Tensor:
"""Continuous actions are usually considered to be independent,.
so we can sum components of the ``log_prob`` or the entropy.
Args:
tensor: shape: (n_batch, n_actions) or (n_batch,)
Returns:
shape: (n_batch,) for (n_batch, n_actions) input, scalar for (n_batch,) input
"""
tensor = tensor.sum(dim=1) if len(tensor.shape) > 1 else tensor.sum()
return tensor
|