Source code for ding.worker.collector.comm.utils

import torch
from ding.interaction.slave import Slave, TaskFail


class NaiveCollector(Slave):
    """
    Overview:
        A slave, whose master is coordinator.
        Used to pass message between comm collector and coordinator.
    Interfaces:
        _process_task, _get_timestep
    """

    def __init__(self, *args, prefix='', **kwargs):
        super().__init__(*args, **kwargs)
        self._prefix = prefix

    def _process_task(self, task):
        """
        Overview:
            Process a task according to input task info dict, which is passed in by master coordinator.
            For each type of task, you can refer to corresponding callback function in comm collector for details.
        Arguments:
            - cfg (:obj:`EasyDict`): Task dict. Must contain key "name".
        Returns:
            - result (:obj:`Union[dict, TaskFail]`): Task result dict, or task fail exception.
        """
        task_name = task['name']
        if task_name == 'resource':
            return {'cpu': '20', 'gpu': '1'}
        elif task_name == 'collector_start_task':
            self.count = 0
            self.task_info = task['task_info']
            return {'message': 'collector task has started'}
        elif task_name == 'collector_data_task':
            self.count += 1
            data_id = './{}_{}_{}'.format(self._prefix, self.task_info['task_id'], self.count)
            torch.save(self._get_timestep(), data_id)
            data = {'data_id': data_id, 'buffer_id': self.task_info['buffer_id'], 'unroll_split_begin': 0}
            data['task_id'] = self.task_info['task_id']
            if self.count == 20:
                return {
                    'task_id': self.task_info['task_id'],
                    'collector_done': True,
                    'cur_episode': 1,
                    'cur_step': 314,
                    'cur_sample': 314,
                }
            else:
                return data
        else:
            raise TaskFail(
                result={'message': 'task name error'}, message='illegal collector task <{}>'.format(task_name)
            )

    def _get_timestep(self):
        return [
            {
                'obs': torch.rand(4),
                'next_obs': torch.randn(4),
                'reward': torch.randint(0, 2, size=(3, )).float(),
                'action': torch.randint(0, 2, size=(1, )),
                'done': False,
            }
        ]