policy

This module contains the CategoricalPolicy implementation.

CategoricalPolicy

A categorical policy parameterized by a neural network.

Source code in src/behavior_generation_lecture_python/mdp/policy.py

class CategoricalPolicy:
    """A categorical policy parameterized by a neural network."""

    def __init__(
        self, sizes: List[int], actions: List[Any], seed: Optional[int] = None
    ) -> None:
        """Initialize the categorical policy.

        Args:
            sizes: List of layer sizes for the MLP.
            actions: List of available actions.
            seed: Random seed for reproducibility (default: None).
        """
        assert sizes[-1] == len(actions)
        if seed is not None:
            torch.manual_seed(seed)
        self.net = multi_layer_perceptron(sizes=sizes)
        self.actions = actions
        self._actions_tensor = torch.tensor(actions, dtype=torch.long).view(
            len(actions), -1
        )

    def _get_distribution(self, state: torch.Tensor) -> Categorical:
        """Calls the model and returns a categorical distribution over the actions.

        Args:
            state: The current state tensor.

        Returns:
            A categorical distribution over actions.
        """
        logits = self.net(state)
        return Categorical(logits=logits)

    def get_action(self, state: torch.Tensor, deterministic: bool = False) -> Any:
        """Returns an action sample for the given state.

        Args:
            state: The current state tensor.
            deterministic: If True, return the most likely action.

        Returns:
            The selected action.
        """
        policy = self._get_distribution(state)
        if deterministic:
            return self.actions[policy.mode.item()]
        return self.actions[policy.sample().item()]

    def get_log_prob(self, states: torch.Tensor, actions: torch.Tensor) -> torch.Tensor:
        """Returns the log-probability for taking the action, when being in the given state.

        Args:
            states: Batch of state tensors.
            actions: Batch of action tensors.

        Returns:
            Log-probabilities of the actions.
        """
        return self._get_distribution(states).log_prob(
            self._get_action_id_from_action(actions)
        )

    def _get_action_id_from_action(self, actions: torch.Tensor) -> torch.Tensor:
        """Returns the indices of the passed actions in self.actions.

        Args:
            actions: Batch of action tensors.

        Returns:
            Tensor of action indices.
        """
        reshaped_actions = actions.unsqueeze(1).expand(
            -1, self._actions_tensor.size(0), -1
        )
        reshaped_actions_tensor = self._actions_tensor.unsqueeze(0).expand(
            actions.size(0), -1, -1
        )
        return torch.where(
            torch.all(reshaped_actions == reshaped_actions_tensor, dim=-1)
        )[1]

__init__(sizes, actions, seed=None)

Initialize the categorical policy.

Parameters:

Name	Type	Description	Default
sizes	List[int]	List of layer sizes for the MLP.	required
actions	List[Any]	List of available actions.	required
seed	Optional[int]	Random seed for reproducibility (default: None).	None

Source code in src/behavior_generation_lecture_python/mdp/policy.py

def __init__(
    self, sizes: List[int], actions: List[Any], seed: Optional[int] = None
) -> None:
    """Initialize the categorical policy.

    Args:
        sizes: List of layer sizes for the MLP.
        actions: List of available actions.
        seed: Random seed for reproducibility (default: None).
    """
    assert sizes[-1] == len(actions)
    if seed is not None:
        torch.manual_seed(seed)
    self.net = multi_layer_perceptron(sizes=sizes)
    self.actions = actions
    self._actions_tensor = torch.tensor(actions, dtype=torch.long).view(
        len(actions), -1
    )

get_action(state, deterministic=False)

Returns an action sample for the given state.

Parameters:

Name	Type	Description	Default
state	Tensor	The current state tensor.	required
deterministic	bool	If True, return the most likely action.	False

Returns:

Type	Description
Any	The selected action.

Source code in src/behavior_generation_lecture_python/mdp/policy.py

def get_action(self, state: torch.Tensor, deterministic: bool = False) -> Any:
    """Returns an action sample for the given state.

    Args:
        state: The current state tensor.
        deterministic: If True, return the most likely action.

    Returns:
        The selected action.
    """
    policy = self._get_distribution(state)
    if deterministic:
        return self.actions[policy.mode.item()]
    return self.actions[policy.sample().item()]

get_log_prob(states, actions)

Returns the log-probability for taking the action, when being in the given state.

Parameters:

Name	Type	Description	Default
states	Tensor	Batch of state tensors.	required
actions	Tensor	Batch of action tensors.	required

Returns:

Type	Description
Tensor	Log-probabilities of the actions.

Source code in src/behavior_generation_lecture_python/mdp/policy.py

def get_log_prob(self, states: torch.Tensor, actions: torch.Tensor) -> torch.Tensor:
    """Returns the log-probability for taking the action, when being in the given state.

    Args:
        states: Batch of state tensors.
        actions: Batch of action tensors.

    Returns:
        Log-probabilities of the actions.
    """
    return self._get_distribution(states).log_prob(
        self._get_action_id_from_action(actions)
    )

multi_layer_perceptron(sizes, activation=nn.ReLU, output_activation=nn.Identity)

Returns a multi-layer perceptron

Source code in src/behavior_generation_lecture_python/mdp/policy.py

def multi_layer_perceptron(
    sizes: List[int],
    activation: Type[nn.Module] = nn.ReLU,
    output_activation: Type[nn.Module] = nn.Identity,
):
    """Returns a multi-layer perceptron"""
    mlp = nn.Sequential()
    for i in range(len(sizes) - 1):
        mlp.append(nn.Linear(sizes[i], sizes[i + 1]))
        if i < len(sizes) - 2:
            mlp.append(activation())
        else:
            mlp.append(output_activation())
    return mlp