viberl.agents.reinforce

REINFORCE: Monte-Carlo policy gradient method for reinforcement learning.

Algorithm Overview:

REINFORCE is a policy gradient method that uses complete episode returns to update the policy parameters. It directly optimizes the policy without requiring a value function, making it conceptually simple but potentially high-variance.

Key Concepts:

• Policy Gradient: Directly optimizes policy parameters \(\theta\) to maximize expected return
• Likelihood Ratio Trick: Uses \(\nabla_\theta \log \pi_\theta(a|s)\) for gradient computation
• Monte-Carlo Returns: Uses complete episode returns \(G_t\) for unbiased estimates
• High Variance: Large variance in gradient estimates but unbiased
• Episode-based Learning: Requires complete episodes before parameter updates

Mathematical Foundation:

Optimization Objective:

\[\nabla_{\theta} J(\theta) = \mathbb{E}\left[\sum_{t=0}^{T} \log \pi_{\theta}(a_t|s_t) G_t\right]\]

Return Calculation:

\[G_t = \sum_{k=0}^{T-t} \gamma^k r_{t+k}\]

Reference: Williams, R.J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning 8, 229-256 (1992). PDF

Classes:

Name	Description
REINFORCEAgent	REINFORCE agent implementation with policy gradient optimization.

REINFORCEAgent

REINFORCEAgent(
    state_size: int,
    action_size: int,
    learning_rate: float = 0.001,
    gamma: float = 0.99,
    hidden_size: int = 128,
    num_hidden_layers: int = 2,
)

Bases: Agent

REINFORCE agent implementation with policy gradient optimization.

This agent implements the REINFORCE algorithm using a policy network to directly optimize the policy parameters via Monte-Carlo gradient estimates.

Parameters:

Name	Type	Description	Default
state_size	int	Dimension of the state space. Must be positive.	required
action_size	int	Number of possible actions. Must be positive.	required
learning_rate	float	Learning rate for the Adam optimizer. Must be positive.	0.001
gamma	float	Discount factor for future rewards. Should be in (0, 1].	0.99
hidden_size	int	Number of neurons in each hidden layer. Must be positive.	128
num_hidden_layers	int	Number of hidden layers. Must be non-negative.	2

Raises:

Type	Description
ValueError	If any parameter is invalid (e.g., negative dimensions).

Initialize the REINFORCE agent.

Methods:

Name	Description
act	Select action using the current policy.
learn	Update policy parameters using the REINFORCE gradient algorithm.
save	Save the agent's policy network parameters to a file.
load	Load the agent's policy network parameters from a file.

Attributes:

Name	Type	Description
gamma
policy_network
optimizer

Source code in viberl/agents/reinforce.py

def __init__(
    self,
    state_size: int,
    action_size: int,
    learning_rate: float = 1e-3,
    gamma: float = 0.99,
    hidden_size: int = 128,
    num_hidden_layers: int = 2,
):
    """Initialize the REINFORCE agent."""
    super().__init__(state_size, action_size)

    if state_size <= 0:
        raise ValueError(f'state_size must be positive, got {state_size}')
    if action_size <= 0:
        raise ValueError(f'action_size must be positive, got {action_size}')
    if learning_rate <= 0:
        raise ValueError(f'learning_rate must be positive, got {learning_rate}')
    if not 0 < gamma <= 1:
        raise ValueError(f'gamma must be in (0, 1], got {gamma}')
    if hidden_size <= 0:
        raise ValueError(f'hidden_size must be positive, got {hidden_size}')
    if num_hidden_layers < 0:
        raise ValueError(f'num_hidden_layers must be non-negative, got {num_hidden_layers}')

    self.gamma = gamma
    self.policy_network = PolicyNetwork(state_size, action_size, hidden_size, num_hidden_layers)
    self.optimizer = optim.Adam(self.policy_network.parameters(), lr=learning_rate)

gamma instance-attribute

gamma = gamma

policy_network instance-attribute

policy_network = PolicyNetwork(state_size, action_size, hidden_size, num_hidden_layers)

optimizer instance-attribute

optimizer = Adam(parameters(), lr=learning_rate)

act

act(state: ndarray, training: bool = True) -> Action

Select action using the current policy.

Uses the policy network to compute action probabilities and selects an action based on the current mode (training vs evaluation).

In training mode, samples from the policy distribution to ensure exploration. In evaluation mode, selects the most probable action (greedy selection).

Parameters:

Name	Type	Description	Default
state	ndarray	Current state observation as a numpy array. Should have shape (state_size,) or be convertible to a tensor of that shape.	required
training	bool	Whether in training mode. If True, samples from the policy distribution for exploration. If False, selects the action with highest probability (greedy selection).	True

Returns:

Type	Description
Action	Action containing the selected action index as an integer.

Raises:

Type	Description
ValueError	If the state has incorrect shape or type.
RuntimeError	If there's an error during action computation.

Source code in viberl/agents/reinforce.py

def act(self, state: np.ndarray, training: bool = True) -> Action:
    """Select action using the current policy.

    Uses the policy network to compute action probabilities and selects
    an action based on the current mode (training vs evaluation).

    In training mode, samples from the policy distribution to ensure
    exploration. In evaluation mode, selects the most probable action
    (greedy selection).

    Args:
        state: Current state observation as a numpy array. Should have shape
            (state_size,) or be convertible to a tensor of that shape.
        training: Whether in training mode. If True, samples from the policy
            distribution for exploration. If False, selects the action with
            highest probability (greedy selection).

    Returns:
        Action containing the selected action index as an integer.

    Raises:
        ValueError: If the state has incorrect shape or type.
        RuntimeError: If there's an error during action computation.
    """
    if training:
        # Training mode: sample from policy distribution
        action = self.policy_network.act(state)
    else:
        # Evaluation mode: select most likely action (greedy)
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            action_probs = self.policy_network(state_tensor)
            action = action_probs.argmax().item()

    return Action(action=action)

learn

learn(trajectories: list[Trajectory]) -> dict[str, float]

Update policy parameters using the REINFORCE gradient algorithm.

Parameters:

Name	Type	Description	Default
trajectories	list[Trajectory]	List of trajectories to learn from	required

Returns:

Type	Description
dict[str, float]	Dictionary containing training metrics:
dict[str, float]	'reinforce/policy_loss': Policy loss value (float)
dict[str, float]	'reinforce/return_mean': Mean of normalized returns (float)
dict[str, float]	'reinforce/batch_size': Number of trajectories in batch (int)

Raises:

Type	Description
ValueError	If no trajectories are provided or contains invalid data.
RuntimeError	If there's an error during gradient computation.

Source code in viberl/agents/reinforce.py

def learn(self, trajectories: list[Trajectory]) -> dict[str, float]:
    """Update policy parameters using the REINFORCE gradient algorithm.

    Args:
        trajectories: List of trajectories to learn from

    Returns:
        Dictionary containing training metrics:
        - 'reinforce/policy_loss': Policy loss value (float)
        - 'reinforce/return_mean': Mean of normalized returns (float)
        - 'reinforce/batch_size': Number of trajectories in batch (int)

    Raises:
        ValueError: If no trajectories are provided or contains invalid data.
        RuntimeError: If there's an error during gradient computation.
    """
    if not trajectories:
        return {}

    # Collect all data from all trajectories
    all_states = []
    all_actions = []
    all_returns = []

    for trajectory in trajectories:
        if not trajectory.transitions:
            continue

        # Extract data from trajectory
        states = [t.state for t in trajectory.transitions]
        actions = [t.action.action for t in trajectory.transitions]
        rewards = [t.reward for t in trajectory.transitions]

        # Compute returns for this trajectory
        returns = self._compute_returns(rewards)

        # Normalize returns for this trajectory
        returns = torch.FloatTensor(returns)
        if returns.std() > 0:
            returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        all_states.extend(states)
        all_actions.extend(actions)
        all_returns.extend(returns.tolist())

    if not all_states:
        return {}

    # Convert to tensors
    states_tensor = torch.FloatTensor(np.array(all_states))
    actions_tensor = torch.LongTensor(all_actions)
    returns_tensor = torch.FloatTensor(all_returns)

    # Get action probabilities
    action_probs = self.policy_network(states_tensor)

    # Compute loss across all trajectories
    m = Categorical(action_probs)
    log_probs = m.log_prob(actions_tensor)
    loss = -torch.mean(log_probs * returns_tensor)

    # Update policy
    self.optimizer.zero_grad()
    loss.backward()
    self.optimizer.step()

    return {
        'reinforce/policy_loss': loss.item(),
        'reinforce/return_mean': returns_tensor.mean().item(),
        'reinforce/batch_size': len(trajectories),
    }

save

save(filepath: str) -> None

Save the agent's policy network parameters to a file.

Saves the complete state of the policy network, including all parameters and buffers. The saved file can be loaded later to restore the exact same policy.

Parameters:

Name	Type	Description	Default
filepath	str	Path where to save the model. Should include the .pth extension. The directory will be created if it doesn't exist.	required

Raises:

Type	Description
IOError	If there's an error writing to the file.
ValueError	If filepath is empty or invalid.

Example

agent.save('models/reinforce_policy.pth')

Later: agent.load('models/reinforce_policy.pth')

Source code in viberl/agents/reinforce.py

def save(self, filepath: str) -> None:
    """Save the agent's policy network parameters to a file.

    Saves the complete state of the policy network, including all parameters
    and buffers. The saved file can be loaded later to restore the exact
    same policy.

    Args:
        filepath: Path where to save the model. Should include the .pth extension.
            The directory will be created if it doesn't exist.

    Raises:
        IOError: If there's an error writing to the file.
        ValueError: If filepath is empty or invalid.

    Example:
        >>> agent.save('models/reinforce_policy.pth')
        >>> # Later: agent.load('models/reinforce_policy.pth')
    """
    if not filepath:
        raise ValueError('filepath cannot be empty')

    # Ensure directory exists
    import os

    os.makedirs(os.path.dirname(filepath), exist_ok=True)

    torch.save(self.policy_network.state_dict(), filepath)

load

load(filepath: str) -> None

Load the agent's policy network parameters from a file.

Loads previously saved policy network parameters from disk. Can be used to restore an agent to a previously saved state.

Parameters:

Name	Type	Description	Default
filepath	str	Path from which to load the model. Should point to a .pth file created by the save() method.	required

Raises:

Type	Description
IOError	If the file doesn't exist or can't be read.
ValueError	If filepath is empty or the file contains invalid data.
RuntimeError	If there's a mismatch between saved and current network architecture.

Example

agent.load('models/reinforce_policy.pth')

Agent is now restored to the saved state

Source code in viberl/agents/reinforce.py

def load(self, filepath: str) -> None:
    """Load the agent's policy network parameters from a file.

    Loads previously saved policy network parameters from disk. Can be used
    to restore an agent to a previously saved state.

    Args:
        filepath: Path from which to load the model. Should point to a .pth file
            created by the save() method.

    Raises:
        IOError: If the file doesn't exist or can't be read.
        ValueError: If filepath is empty or the file contains invalid data.
        RuntimeError: If there's a mismatch between saved and current network architecture.

    Example:
        >>> agent.load('models/reinforce_policy.pth')
        >>> # Agent is now restored to the saved state
    """
    if not filepath:
        raise ValueError('filepath cannot be empty')

    if not os.path.exists(filepath):
        raise OSError(f'File not found: {filepath}')

    state_dict = torch.load(filepath, map_location='cpu')

    # Verify compatibility
    current_params = self.policy_network.state_dict()
    if state_dict.keys() != current_params.keys():
        raise RuntimeError(
            'Network architecture mismatch: saved model has different parameters'
        )

    self.policy_network.load_state_dict(state_dict)