ml-agents/ml-agents/mlagents/trainers/tests/test_simple_transfer.py

import math
import tempfile
import pytest
import numpy as np
import attr
from typing import Dict

from mlagents.trainers.tests.transfer_test_envs import SimpleTransferEnvironment
from mlagents.trainers.trainer_controller import TrainerController
from mlagents.trainers.trainer_util import TrainerFactory
from mlagents.trainers.simple_env_manager import SimpleEnvManager
from mlagents.trainers.demo_loader import write_demo
from mlagents.trainers.stats import StatsReporter, StatsWriter, StatsSummary, TensorboardWriter, CSVWriter
from mlagents.trainers.settings import (
    TrainerSettings,
    PPOSettings,
    PPOTransferSettings,
    SACSettings,
    NetworkSettings,
    SelfPlaySettings,
    BehavioralCloningSettings,
    GAILSettings,
    TrainerType,
    RewardSignalType,
)
from mlagents.trainers.models import EncoderType, ScheduleType
from mlagents_envs.side_channel.environment_parameters_channel import (
    EnvironmentParametersChannel,
)
from mlagents_envs.communicator_objects.demonstration_meta_pb2 import (
    DemonstrationMetaProto,
)
from mlagents_envs.communicator_objects.brain_parameters_pb2 import BrainParametersProto
from mlagents_envs.communicator_objects.space_type_pb2 import discrete, continuous

BRAIN_NAME = "Simple"


PPO_CONFIG = TrainerSettings(
    trainer_type=TrainerType.PPO,
    hyperparameters=PPOSettings(
        learning_rate=5.0e-3,
        learning_rate_schedule=ScheduleType.CONSTANT,
        batch_size=16,
        buffer_size=64,
    ),
    network_settings=NetworkSettings(num_layers=2, hidden_units=32),
    summary_freq=500,
    max_steps=3000,
    threaded=False,
)

SAC_CONFIG = TrainerSettings(
    trainer_type=TrainerType.SAC,
    hyperparameters=SACSettings(
        learning_rate=5.0e-3,
        learning_rate_schedule=ScheduleType.CONSTANT,
        batch_size=8,
        buffer_init_steps=100,
        buffer_size=5000,
        tau=0.01,
        init_entcoef=0.01,
    ),
    network_settings=NetworkSettings(num_layers=1, hidden_units=16),
    summary_freq=100,
    max_steps=1000,
    threaded=False,
)

Transfer_CONFIG = TrainerSettings(
    trainer_type=TrainerType.PPO_Transfer,
    hyperparameters=PPOTransferSettings(
        learning_rate=5.0e-3,
        learning_rate_schedule=ScheduleType.CONSTANT,
        batch_size=16,
        buffer_size=64,
        feature_size=4,
        reuse_encoder=True,
        in_epoch_alter=True,
        # in_batch_alter=True,
        use_op_buffer=True,
        # policy_layers=0,
        # value_layers=0,
        # conv_thres=1e-4,
        # predict_return=True
        # separate_policy_train=True,
        # separate_value_train=True
        # separate_value_net=True,
    ),
    network_settings=NetworkSettings(num_layers=1, hidden_units=32),
    summary_freq=500,
    max_steps=3000,
    threaded=False,
)



# The reward processor is passed as an argument to _check_environment_trains.
# It is applied to the list pf all final rewards for each brain individually.
# This is so that we can process all final rewards in different ways for different algorithms.
# Custom reward processors shuld be built within the test function and passed to _check_environment_trains
# Default is average over the last 5 final rewards
def default_reward_processor(rewards, last_n_rewards=5):
    rewards_to_use = rewards[-last_n_rewards:]
    # For debugging tests
    print("Last {} rewards:".format(last_n_rewards), rewards_to_use)
    return np.array(rewards[-last_n_rewards:], dtype=np.float32).mean()


class DebugWriter(StatsWriter):
    """
    Print to stdout so stats can be viewed in pytest
    """

    def __init__(self):
        self._last_reward_summary: Dict[str, float] = {}
        self.stats = {}

    def get_last_rewards(self):
        return self._last_reward_summary

    def write_stats(
        self, category: str, values: Dict[str, StatsSummary], step: int
    ) -> None:
        for val, stats_summary in values.items():
            if val == "Environment/Cumulative Reward":
                print(step, val, stats_summary.mean)
                self.stats[step] = stats_summary.mean
                self._last_reward_summary[category] = stats_summary.mean
    
    def write2file(self, filename):
        with open(filename, "w") as reward_file:
            for step in self.stats.keys():
                reward_file.write(str(step) + ":" + str(self.stats[step]) + "\n")


def _check_environment_trains(
    env,
    trainer_config,
    reward_processor=default_reward_processor,
    meta_curriculum=None,
    success_threshold=0.9,
    env_manager=None,
    run_id="id",
    seed=1337
):
    # Create controller and begin training.
    model_dir = "./transfer_results/" + run_id
    StatsReporter.writers.clear()  # Clear StatsReporters so we don't write to file
    debug_writer = DebugWriter()
    StatsReporter.add_writer(debug_writer)

    csv_writer = CSVWriter(
        model_dir,
        required_fields=[
            "Environment/Cumulative Reward",
            "Environment/Episode Length",
        ],
    )
    tb_writer = TensorboardWriter(
        model_dir, clear_past_data=True
    )
    StatsReporter.add_writer(tb_writer)
    StatsReporter.add_writer(csv_writer)

    if env_manager is None:
        env_manager = SimpleEnvManager(env, EnvironmentParametersChannel())
    trainer_factory = TrainerFactory(
        trainer_config=trainer_config,
        output_path=model_dir,
        train_model=True,
        load_model=False,
        seed=seed,
        meta_curriculum=meta_curriculum,
        multi_gpu=False,
    )

    tc = TrainerController(
        trainer_factory=trainer_factory,
        output_path=model_dir,
        run_id=run_id,
        meta_curriculum=meta_curriculum,
        train=True,
        training_seed=seed,
    )

    # Begin training
    tc.start_learning(env_manager)
    # debug_writer.write2file(model_dir+"/reward.txt")

    # if (
    #     success_threshold is not None
    # ):  # For tests where we are just checking setup and not reward
    #     processed_rewards = [
    #         reward_processor(rewards) for rewards in env.final_rewards.values()
    #     ]
    #     assert all(not math.isnan(reward) for reward in processed_rewards)
    #     assert all(reward > success_threshold for reward in processed_rewards)


def test_2d_model(config=Transfer_CONFIG, obs_spec_type="normal", run_id="model_normal_f4_varp-pri-test", seed=0):
    env = SimpleTransferEnvironment(
        [BRAIN_NAME], use_discrete=False, action_size=2, step_size=0.1, 
        num_vector=2, obs_spec_type=obs_spec_type, goal_type="hard"
    )
    new_hyperparams = attr.evolve(
        config.hyperparameters, batch_size=120, buffer_size=12000, learning_rate=5.0e-3,
        use_var_predict=True, with_prior=True
    )
    config = attr.evolve(config, hyperparameters=new_hyperparams, max_steps=200000, summary_freq=5000)
    _check_environment_trains(env, {BRAIN_NAME: config}, run_id=run_id + "_s" + str(seed), seed=seed)

def test_2d_transfer(config=Transfer_CONFIG, obs_spec_type="rich1", run_id="transfer_rich1_from-normal_varp-pri_retrain-all_5e-2", seed=1337):
    env = SimpleTransferEnvironment(
        [BRAIN_NAME], use_discrete=False, action_size=2, step_size=0.1, 
        num_vector=2, obs_spec_type=obs_spec_type, goal_type="hard"
    )
    new_hyperparams = attr.evolve(
        config.hyperparameters, batch_size=120, buffer_size=12000, use_transfer=True,
        transfer_path="./transfer_results/model_normal_f4_varp-pri_s0/Simple",
        use_op_buffer=True, in_epoch_alter=True, in_batch_alter=False, learning_rate=5.0e-2, 
        train_policy=True, train_value=True, train_model=True, feature_size=4, learning_rate_schedule=ScheduleType.LINEAR,
        use_var_predict=True, with_prior=True
    )
    config = attr.evolve(config, hyperparameters=new_hyperparams, max_steps=200000, summary_freq=5000)
    _check_environment_trains(env, {BRAIN_NAME: config}, run_id=run_id + "_s" + str(seed), seed=seed)


if __name__ == "__main__":
    # test_2d_model(seed=0)
    # test_2d_model(config=PPO_CONFIG, run_id="ppo_normal", seed=0)
    test_2d_transfer(seed=123)
    # for i in range(5):
    #     test_2d_model(seed=i)