<filename>dedo/utils/rl_sb3_utils.py<gh_stars>10-100
"""
Utilities for RL training and eval.


Note: this code is for research i.e. quick experimentation; it has minimal
comments for now, but if we see further interest from the community -- we will
add further comments, unify the style, improve efficiency and add unittests.

@contactrika

"""
import os
import pickle

import cv2
import torch

from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.logger import Video

from dedo.utils.train_utils import object_to_str


def play(env, num_episodes, rl_agent, debug=False, logdir=None,
         cam_resolution=None, filename="play"):
    vidwriter = None
    if logdir is not None:
        if not os.path.exists(logdir):
            os.mkdir(logdir)
        vidwriter = cv2.VideoWriter(
            os.path.join(logdir, f'{filename}.mp4'), cv2.VideoWriter_fourcc(*'mp4v'),
            24, (cam_resolution, cam_resolution))
    for epsd in range(num_episodes):
        episode_rwd = 0.0
        if debug:
            print('------------ Play episode ', epsd, '------------------')
        obs = env.reset()
        if vidwriter is not None:
            img = env.render(mode='rgb_array', width=cam_resolution,
                             height=cam_resolution)
            vidwriter.write(img[..., ::-1])
        step = 0
        while True:
            # rl_agent.predict() to get acts, not forcing deterministic.
            act, _states = rl_agent.predict(obs)
            next_obs, rwd, done, info = env.step(act)
            episode_rwd += rwd
            if vidwriter is not None:
                img = env.render(mode='rgb_array', width=cam_resolution,
                                 height=cam_resolution)
                vidwriter.write(img[..., ::-1])
            if done:
                if debug:
                    print(f'episode reward {episode_rwd:0.2f}')
                break
            obs = next_obs
            step += 1
        # input('Episode ended; press enter to go on')
    if vidwriter is not None:
        vidwriter.release()


class CustomCallback(BaseCallback):
    """
    A custom callback that runs eval and adds videos to Tensorboard.
    """

    def __init__(self, eval_env, logdir, num_train_envs, args,
                 num_steps_between_save=10000, viz=False, debug=False):
        super(CustomCallback, self).__init__(debug)
        # Those variables will be accessible in the callback
        # (they are defined in the base class)
        # The RL model
        # self.model = None  # type: BaseAlgorithm
        # An alias for self.model.get_env(), the environment used for training
        # self.training_env = None  # type: Union[gym.Env, VecEnv, None]
        # Number of time the callback was called
        # self.n_calls = 0  # type: int
        # self.num_timesteps = 0  # type: int
        # local and global variables
        # self.locals = None  # type: Dict[str, Any]
        # self.globals = None  # type: Dict[str, Any]
        # The logger object, used to report things in the terminal
        # self.logger = None  # stable_baselines3.common.logger
        # # Sometimes, for event callback, it is useful
        # # to have access to the parent object
        # self.parent = None  # type: Optional[BaseCallback]
        self._eval_env = eval_env
        self._logdir = logdir
        self._num_train_envs = num_train_envs
        self._my_args = args
        self._num_steps_between_save = num_steps_between_save
        self._viz = viz
        self._debug = debug
        self._steps_since_save = num_steps_between_save  # save right away

    def _on_training_start(self) -> None:
        """
        This method is called before the first rollout starts.
        """
        # Log args to tensorboard.
        self.logger.record('args', object_to_str(self._my_args))

    def _on_rollout_start(self) -> None:
        """
        A rollout is the collection of environment interaction
        using the current policy.
        This event is triggered before collecting new samples.
        """
        pass

    def _on_step(self) -> bool:
        """
        This method will be called by the model after each call to `env.step()`.

        For child callback (of an `EventCallback`), this will be called
        when the event is triggered.

        :return: (bool) If the callback returns False, training is aborted early.
        """
        self._steps_since_save += self._num_train_envs
        if self._steps_since_save >= self._num_steps_between_save:
            # Save checkpoint.
            if self._logdir is not None:
                self.model.save(os.path.join(self._logdir, 'agent'))
                pickle.dump(self._my_args,
                            open(os.path.join(self._logdir, 'args.pkl'), 'wb'),
                            protocol=pickle.HIGHEST_PROTOCOL)
            self._steps_since_save = 0
            # Record video.
            if not self._my_args.disable_logging_video:
                screens = []

                def grab_screens(_locals, _globals):
                    screen = self._eval_env.render(
                        mode='rgb_array', width=300, height=300)
                    # PyTorch uses CxHxW vs HxWxC gym (and TF) images
                    screens.append(screen.transpose(2, 0, 1))

                evaluate_policy(
                    self.model, self._eval_env, callback=grab_screens,
                    n_eval_episodes=1, deterministic=False)
                self.logger.record(
                    'trajectory/video',
                    Video(torch.ByteTensor([screens]), fps=50),
                    exclude=('stdout', 'log', 'json', 'csv'))

        return True

    def _on_rollout_end(self) -> None:
        """
        This event is triggered before updating the policy.
        """
        pass

    def _on_training_end(self) -> None:
        """
        This event is triggered before exiting the `learn()` method.
        """
        pass

<filename>dedo/utils/rl_sb3_utils.py<gh_stars>10-100
"""
Utilities for RL training and eval.


Note: this code is for research i.e. quick experimentation; it has minimal
comments for now, but if we see further interest from the community -- we will
add further comments, unify the style, improve efficiency and add unittests.

@contactrika

"""
import os
import pickle

import cv2
import torch

from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.logger import Video

from dedo.utils.train_utils import object_to_str


def play(env, num_episodes, rl_agent, debug=False, logdir=None,
         cam_resolution=None, filename="play"):
    vidwriter = None
    if logdir is not None:
        if not os.path.exists(logdir):
            os.mkdir(logdir)
        vidwriter = cv2.VideoWriter(
            os.path.join(logdir, f'{filename}.mp4'), cv2.VideoWriter_fourcc(*'mp4v'),
            24, (cam_resolution, cam_resolution))
    for epsd in range(num_episodes):
        episode_rwd = 0.0
        if debug:
            print('------------ Play episode ', epsd, '------------------')
        obs = env.reset()
        if vidwriter is not None:
            img = env.render(mode='rgb_array', width=cam_resolution,
                             height=cam_resolution)
            vidwriter.write(img[..., ::-1])
        step = 0
        while True:
            # rl_agent.predict() to get acts, not forcing deterministic.
            act, _states = rl_agent.predict(obs)
            next_obs, rwd, done, info = env.step(act)
            episode_rwd += rwd
            if vidwriter is not None:
                img = env.render(mode='rgb_array', width=cam_resolution,
                                 height=cam_resolution)
                vidwriter.write(img[..., ::-1])
            if done:
                if debug:
                    print(f'episode reward {episode_rwd:0.2f}')
                break
            obs = next_obs
            step += 1
        # input('Episode ended; press enter to go on')
    if vidwriter is not None:
        vidwriter.release()


class CustomCallback(BaseCallback):
    """
    A custom callback that runs eval and adds videos to Tensorboard.
    """

    def __init__(self, eval_env, logdir, num_train_envs, args,
                 num_steps_between_save=10000, viz=False, debug=False):
        super(CustomCallback, self).__init__(debug)
        # Those variables will be accessible in the callback
        # (they are defined in the base class)
        # The RL model
        # self.model = None  # type: BaseAlgorithm
        # An alias for self.model.get_env(), the environment used for training
        # self.training_env = None  # type: Union[gym.Env, VecEnv, None]
        # Number of time the callback was called
        # self.n_calls = 0  # type: int
        # self.num_timesteps = 0  # type: int
        # local and global variables
        # self.locals = None  # type: Dict[str, Any]
        # self.globals = None  # type: Dict[str, Any]
        # The logger object, used to report things in the terminal
        # self.logger = None  # stable_baselines3.common.logger
        # # Sometimes, for event callback, it is useful
        # # to have access to the parent object
        # self.parent = None  # type: Optional[BaseCallback]
        self._eval_env = eval_env
        self._logdir = logdir
        self._num_train_envs = num_train_envs
        self._my_args = args
        self._num_steps_between_save = num_steps_between_save
        self._viz = viz
        self._debug = debug
        self._steps_since_save = num_steps_between_save  # save right away

    def _on_training_start(self) -> None:
        """
        This method is called before the first rollout starts.
        """
        # Log args to tensorboard.
        self.logger.record('args', object_to_str(self._my_args))

    def _on_rollout_start(self) -> None:
        """
        A rollout is the collection of environment interaction
        using the current policy.
        This event is triggered before collecting new samples.
        """
        pass

    def _on_step(self) -> bool:
        """
        This method will be called by the model after each call to `env.step()`.

        For child callback (of an `EventCallback`), this will be called
        when the event is triggered.

        :return: (bool) If the callback returns False, training is aborted early.
        """
        self._steps_since_save += self._num_train_envs
        if self._steps_since_save >= self._num_steps_between_save:
            # Save checkpoint.
            if self._logdir is not None:
                self.model.save(os.path.join(self._logdir, 'agent'))
                pickle.dump(self._my_args,
                            open(os.path.join(self._logdir, 'args.pkl'), 'wb'),
                            protocol=pickle.HIGHEST_PROTOCOL)
            self._steps_since_save = 0
            # Record video.
            if not self._my_args.disable_logging_video:
                screens = []

                def grab_screens(_locals, _globals):
                    screen = self._eval_env.render(
                        mode='rgb_array', width=300, height=300)
                    # PyTorch uses CxHxW vs HxWxC gym (and TF) images
                    screens.append(screen.transpose(2, 0, 1))

                evaluate_policy(
                    self.model, self._eval_env, callback=grab_screens,
                    n_eval_episodes=1, deterministic=False)
                self.logger.record(
                    'trajectory/video',
                    Video(torch.ByteTensor([screens]), fps=50),
                    exclude=('stdout', 'log', 'json', 'csv'))

        return True

    def _on_rollout_end(self) -> None:
        """
        This event is triggered before updating the policy.
        """
        pass

    def _on_training_end(self) -> None:
        """
        This event is triggered before exiting the `learn()` method.
        """
        pass

Utilities for RL training and eval.   Note: this code is for research i.e. quick experimentation; it has minimal comments for now, but if we see further interest from the community -- we will add further comments, unify the style, improve efficiency and add unittests.  @contactrika # rl_agent.predict() to get acts, not forcing deterministic. # input('Episode ended; press enter to go on') A custom callback that runs eval and adds videos to Tensorboard. # Those variables will be accessible in the callback # (they are defined in the base class) # The RL model # self.model = None  # type: BaseAlgorithm # An alias for self.model.get_env(), the environment used for training # self.training_env = None  # type: Union[gym.Env, VecEnv, None] # Number of time the callback was called # self.n_calls = 0  # type: int # self.num_timesteps = 0  # type: int # local and global variables # self.locals = None  # type: Dict[str, Any] # self.globals = None  # type: Dict[str, Any] # The logger object, used to report things in the terminal # self.logger = None  # stable_baselines3.common.logger # # Sometimes, for event callback, it is useful # # to have access to the parent object # self.parent = None  # type: Optional[BaseCallback] # save right away This method is called before the first rollout starts. # Log args to tensorboard. A rollout is the collection of environment interaction         using the current policy.         This event is triggered before collecting new samples. This method will be called by the model after each call to `env.step()`.          For child callback (of an `EventCallback`), this will be called         when the event is triggered.          :return: (bool) If the callback returns False, training is aborted early. # Save checkpoint. # Record video. # PyTorch uses CxHxW vs HxWxC gym (and TF) images This event is triggered before updating the policy. This event is triggered before exiting the `learn()` method.