ml-agents/ml-agents/mlagents/trainers/optimizer/torch_optimizer.py

from typing import Dict, Optional, Tuple, List
from mlagents.torch_utils import torch
import numpy as np
import math

from mlagents.trainers.buffer import AgentBuffer, AgentBufferField
from mlagents.trainers.trajectory import ObsUtil
from mlagents.trainers.torch.components.bc.module import BCModule
from mlagents.trainers.torch.components.reward_providers import create_reward_provider

from mlagents.trainers.policy.torch_policy import TorchPolicy
from mlagents.trainers.optimizer import Optimizer
from mlagents.trainers.settings import TrainerSettings
from mlagents.trainers.torch.utils import ModelUtils


class TorchOptimizer(Optimizer):
    def __init__(self, policy: TorchPolicy, trainer_settings: TrainerSettings):
        super().__init__()
        self.policy = policy
        self.trainer_settings = trainer_settings
        self.update_dict: Dict[str, torch.Tensor] = {}
        self.value_heads: Dict[str, torch.Tensor] = {}
        self.memory_in: torch.Tensor = None
        self.memory_out: torch.Tensor = None
        self.m_size: int = 0
        self.global_step = torch.tensor(0)
        self.bc_module: Optional[BCModule] = None
        self.create_reward_signals(trainer_settings.reward_signals)
        self.critic_memory_dict: Dict[str, torch.Tensor] = {}
        if trainer_settings.behavioral_cloning is not None:
            self.bc_module = BCModule(
                self.policy,
                trainer_settings.behavioral_cloning,
                policy_learning_rate=trainer_settings.hyperparameters.learning_rate,
                default_batch_size=trainer_settings.hyperparameters.batch_size,
                default_num_epoch=3,
            )

    @property
    def critic(self):
        raise NotImplementedError

    def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
        pass

    def create_reward_signals(self, reward_signal_configs):
        """
        Create reward signals
        :param reward_signal_configs: Reward signal config.
        """
        for reward_signal, settings in reward_signal_configs.items():
            # Name reward signals by string in case we have duplicates later
            self.reward_signals[reward_signal.value] = create_reward_provider(
                reward_signal, self.policy.behavior_spec, settings
            )

    def _evaluate_by_sequence(
        self, tensor_obs: List[torch.Tensor], initial_memory: np.ndarray
    ) -> Tuple[Dict[str, torch.Tensor], AgentBufferField, torch.Tensor]:
        """
        Evaluate the batch sequence-by-sequence, assembling the result. This enables us to get the
        intermediate memories for the critic.
        """
        num_experiences = tensor_obs[0].shape[0]
        all_next_memories = AgentBufferField()
        # The 1st sequence are the ones that are padded. So if seq_len = 3 and
        # trajectory is of length 10, the ist sequence is [pad,pad,obs].
        # Compute the number of elements in this padded seq.
        leftover = num_experiences % self.policy.sequence_length
        first_seq_len = self.policy.sequence_length if leftover == 0 else leftover
        for _ in range(first_seq_len):
            all_next_memories.append(initial_memory.squeeze().detach().numpy())

        # Compute values for the potentially truncated initial sequence

        _mem = initial_memory
        seq_obs = []
        for _obs in tensor_obs:
            if leftover > 0:
                # Pad
                padding_shape = list(_obs.shape)
                padding_shape[0] = self.policy.sequence_length - leftover
                padding = torch.zeros(padding_shape)
                padded_obs = torch.cat([padding, _obs[0:leftover]])
            else:
                padded_obs = _obs[0 : self.policy.sequence_length]
            seq_obs.append(padded_obs)
        init_values, _mem = self.critic.critic_pass(
            seq_obs, _mem, sequence_length=self.policy.sequence_length
        )
        # Trim out padded part, i.e. get last leftover number of elements
        all_values = {
            signal_name: [init_values[signal_name][-leftover:]]
            for signal_name in init_values.keys()
        }

        # Evaluate other trajectories
        for seq_num in range(
            1, math.ceil((num_experiences) / (self.policy.sequence_length))
        ):
            seq_obs = []
            for _obs in tensor_obs:
                start = seq_num * self.policy.sequence_length - (
                    self.policy.sequence_length - leftover
                )
                end = (seq_num + 1) * self.policy.sequence_length - (
                    self.policy.sequence_length - leftover
                )
                seq_obs.append(_obs[start:end])
            values, _mem = self.critic.critic_pass(
                seq_obs, _mem, sequence_length=self.policy.sequence_length
            )
            for _ in range(self.policy.sequence_length):
                all_next_memories.append(_mem.squeeze().detach().numpy())
            for signal_name, _val in values.items():
                all_values[signal_name].append(_val)

        # Create one tensor per reward signal
        all_value_tensors = {
            signal_name: torch.cat(value_list, dim=0)
            for signal_name, value_list in all_values.items()
        }
        next_mem = _mem
        return all_value_tensors, all_next_memories, next_mem

    def get_trajectory_value_estimates(
        self,
        batch: AgentBuffer,
        next_obs: List[np.ndarray],
        done: bool,
        agent_id: str = "",
    ) -> Tuple[Dict[str, np.ndarray], Dict[str, float], Optional[AgentBufferField]]:
        n_obs = len(self.policy.behavior_spec.observation_specs)

        if agent_id in self.critic_memory_dict:
            memory = self.critic_memory_dict[agent_id]
        else:
            memory = (
                torch.zeros((1, 1, self.critic.memory_size))
                if self.policy.use_recurrent
                else None
            )

        # Convert to tensors
        current_obs = ObsUtil.from_buffer(batch, n_obs)
        current_obs = [ModelUtils.list_to_tensor(obs) for obs in current_obs]
        next_obs = [ModelUtils.list_to_tensor(obs) for obs in next_obs]

        next_obs = [obs.unsqueeze(0) for obs in next_obs]

        # If we're using LSTM, we want to get all the intermediate memories.
        all_next_memories: Optional[AgentBufferField] = None
        if self.policy.use_recurrent:
            (
                value_estimates,
                all_next_memories,
                next_memory,
            ) = self._evaluate_by_sequence(current_obs, memory)
        else:
            value_estimates, next_memory = self.critic.critic_pass(
                current_obs, memory, sequence_length=batch.num_experiences
            )

        # Store the memory for the next trajectory
        self.critic_memory_dict[agent_id] = next_memory

        next_value_estimate, _ = self.critic.critic_pass(
            next_obs, next_memory, sequence_length=1
        )

        for name, estimate in value_estimates.items():
            value_estimates[name] = ModelUtils.to_numpy(estimate)
            next_value_estimate[name] = ModelUtils.to_numpy(next_value_estimate[name])

        if done:
            for k in next_value_estimate:
                if not self.reward_signals[k].ignore_done:
                    next_value_estimate[k] = 0.0
            if agent_id in self.critic_memory_dict:
                self.critic_memory_dict.pop(agent_id)
        return value_estimates, next_value_estimate, all_next_memories