Rllib学习[2] --env定义 + env rollout

本文讲述rllib中环境如何搭建，以及环境中sample模块。

env

目前还没有涉及到VectorEnv，主要使用gym.Env 以及 基于gym.Env创建的子类环境。

为了在ray 和 rllib中能够使用 自己创建的环境，需要注意：
初始化中含有 observation_space 和 action_space 用于rllib后续policy中建立神经网络的输入。

基本框架

# env需要按照 gym.env 进行设置
class MyEnv(gym.Env):
    def __init__(self, env_config):  # 此处的变量请只使用 dict env_config 用于定义所有的参数
    								# 除非使用register建立，同时需要在 env_creator中合理对应
        self.action_space = <gym.Space>
        self.observation_space = <gym.Space>
    def reset(self):
        return <obs>
    def step(self, action):
        return <obs>, <reward: float>, <done: bool>, <info: dict>

gym.observation_space 配置

环境的调用过程

import gym, ray
from ray.rllib.algorithms import ppo

ray.init()
algo = ppo.PPOTrainer(env=MyEnv, config={
    "env_config": {},  # config to pass to env class
})

环境的注册

我们可以使用register（注册），将我们自己定义的环境设置成 rllib可以识别的 环境string，然后可以直接进行调用。
需要注意的是： gym中的registry和 ray不完全兼容。因此请使用ray中的resgister进行注册。

from ray.tune.registry import register_env

def env_creator(env_config):   # 此处的 env_config对应 我们在建立trainer时传入的dict env_config
    return MyEnv(...)  # return an env instance

register_env("my_env", env_creator) # 此处传入了 环境的名称 | 环境的实例调用函数
algo = ppo.PPO(env="my_env",config={
									"env_config":{} # 传入 env的__init__中
})

多环境的建立

有时候一个worker需要对多个环境进行学习，因此我们在定义环境的时候 需要同时定义几个环境。

对于 num_envs_per_worker >0 的情况， 每一个worker会对应有多个环境，因此需要通过根据env_config.worker_index 和 env_config.vector_index 来得到 worker的 id和 env id

class MultiEnv(gym.Env):
    def __init__(self, env_config):
        # pick actual env based on worker and env indexes
        self.env = gym.make(
            choose_env_for(env_config.worker_index, env_config.vector_index))
        self.action_space = self.env.action_space
        self.observation_space = self.env.observation_space
    def reset(self):
        return self.env.reset()
    def step(self, action):
        return self.env.step(action)

register_env("multienv", lambda config: MultiEnv(config))

注意事项

如果我们需要在环境中进行log，我们需要将log的设置也设置在环境中，然后会在worker内运行。
和 callback中的log如何结合？

sacling guide

多智能体+分层

trainer sampling

worker是包含于trainer当中的， workerset是 一个 local_worker 和 多个(num_workers) remote_worker，每一个worker都是一个 RolloutWorker。而每一个worker会进行一个或多个的env训练。 每个worker生成的数据都会形成 SampleBatch 的 类，这个类包含了训练需要的数据。不同worker生成的SampleBatch都可以整合。

WorkerSet

在 trainer类中 会有 trainer.workers，表示所有的RollWorker 形成的 WorkSet。 想要调用每一个worker时， local worker则是 trainer.workers.local_worker()， 对于remote_worker 则是 trainer.workers.remote_workers() 将返回 一个list，包含所有的 remote_workers。
对于evaluation，会生成单独的WorkerSet: self.evaluation_workers.

import ray.rllib.evaluation.worker_set.WorkerSet
class ray.rllib.evaluation.worker_set.WorkerSet(*, env_creator: Optional[Callable[[ray.rllib.env.env_context.EnvContext], Optional[Any]]] = None, validate_env: Optional[Callable[[Any], None]] = None, policy_class: Optional[Type[ray.rllib.policy.policy.Policy]] = None, trainer_config: Optional[dict] = None, num_workers: int = 0, local_worker: bool = True, logdir: Optional[str] = None, _setup: bool = True):

def sync_weights(policies: Optional[List[str]] = None, from_worker: Optional[ray.rllib.evaluation.rollout_worker.RolloutWorker] = None, global_vars: Optional[Dict[str, Any]] = None)-> None 
# policies 表示更新的policyID ； from_worker 表示更新时 policy的来源，默认是local_worker

def add_workers(num_workers: int)-> None
# 在WorkerSet中 增加 remote_workers
def foreach_worker(func) 
# func(worker) 以WorkerSet中的每一个worker为参数，生成值。 foreach_worker会对每一个worker进行运算，返回 所有的值

def foreach_worker_with_index(func)-> list
# func(worker,index) 将 worker 和 对应的 index 作为输入进行操作。

def foreach_policy(func) -> list
# func(policy,policy_id) 以每个worker中的policy和id为参数。 对于多智能体，会出现多个id；单智能体只会有一个id。

RolloutWorker

 # 会包装一个policy实例 和一个环境类来收集经验
class ray.rllib.evaluation.rollout_worker.RolloutWorker(*, 
		env_creator,  # 创建环境的func
		validate_env = None, # 测试env的函数
		policy_spec = Union[type,dict], #单policy：  policy ； 				 
											# 多policy： {"policy1":policy1,"policy2":policy2}
		policy_mapping_fn = Optional[Callable[[Any, Episode], str]], # 多智能体中 将agent与policy对应
		rollout_fragment_length: int , # 每次训练的pari数目
		batch_mode: str, # 默认"truncated_episodes" ; 也可以"complete_episodes"
		episode_horizon: int, # 每个episode的最大步长
		preprocessor_pref: str,  # 是否使用 preprocessor, 默认是deepmind，可以使用rllib
		sample_async : bool, # 是否使用非同步方式收集sample， 如果是的话 收集信息会快，但是会有些 off-policy
		compress_observations: bool, # 是否压缩 obs
		num_envs: int, # 每一个worker 采用的env 多少
		clip_rewards: bool , # 是否截断reward
		normalize_actions: bool, # 是否将action正则化到 指定范围内
		clip_actions: bool, # 是否将动作截断 
		env_config, # 传入到 env_config中
		model_config， # 用于创建 policy model
		policy_config，# 传入policy中。 多智能体中会和policy_mapping一起使用
		worker_index, # 设置当前RollWorker的 index， 对于remote worker 需要设置index>0.
		num_workers, # 用于设置remote worker的数量。 对于单worker时，num_workers应该设置为0
		recreated_worker: bool, # 用于查看 该RollWorker是否是一开始创建的，还是worker fail之后重新创建的
								# 对于一开始创建的worker 和 后创建的worker 只有这一项不同
		record_env: bool, #  是否记录 每个episode的数据和video
		log_dir: str, # log记录的位置， 默认是~/ray_results
		callbacks: class, # 自定义的回调函数
        soft_horizon: bool = False, # 到了horizon 不重置，继续计算reward
        no_done_at_end: bool = False,  # 不到done 不结束
        seed: int = None, # 随机种子
        spaces: Optional[Dict[PolicyID, Tuple[Space, Space]]], 
        # 在没有env创建时，生成policy id到 obs_space,action_space 的 dict映射。 
		disable_env_checking=False, # 用于不检查 env
	)
	def sample(self)  # 会返回 BatchSample，生成需要的sample
	# 生成的具体信息 详见 下节的 BatchSample

	def sample_with_count(self) # 返回 (BatchSample, count_number)  count_number 表示 BatchSample中的pair的数量
	
	def learn_on_batch(samples) # input是BatchSample， 然后更新policy

	def sample_and_learn(expected_batch_size: int, num_sgd_iter: int, sgd_minibatch_size: str, standardize_fields: List[str]) -> Tuple[dict, int]
	# expected_batch_size : batch的大小 ； num_sgd_iter ： 迭代的次数； sgd_minibatch_size： 每次更新的mini-batch size ； standardize_fields: samples 中 需要正则化的项 
	# 实际上是调用 policy

	def compute_gradients(samples) # 计算samples对应的gradient
	
	def apply_gradients(grads) # 根据grads 进行更新

	def get_metrics() # 得到对应的metrics
 	def foreach_env(func: Callable[[Any], ray.rllib.utils.typing.T]) → List[ray.rllib.utils.typing.T]
		# func 将会对每一个env进行
		
	def get_policy(policy_id: str = 'default_policy') # 根据policy id得到对应的policy
	
	def save()
	def restore()
	def get_weights() # 得到model的weights
	def set_weights() # 设置model的weights

初始化RolloutWorker

import gym 
from ray.rllib.evaluation.rollout_worker import RolloutWorker
from ray.rllib.agents.pg.pg_tf_policy import PGTFPolicy
worker = RolloutWorker( 
  env_creator=lambda _: gym.make("CartPole-v0"), 
  policy_spec=PGTFPolicy)

SampleBatch

{ 'action_logp': np.ndarray((200,), dtype=float32, min=-0.701, max=-0.685, mean=-0.694),
  'actions': np.ndarray((200,), dtype=int64, min=0.0, max=1.0, mean=0.495),
  'dones': np.ndarray((200,), dtype=bool, min=0.0, max=1.0, mean=0.055),
  'infos': np.ndarray((200,), dtype=object, head={}),
  'new_obs': np.ndarray((200, 4), dtype=float32, min=-2.46, max=2.259, mean=0.018),
  'obs': np.ndarray((200, 4), dtype=float32, min=-2.46, max=2.259, mean=0.016),
  'rewards': np.ndarray((200,), dtype=float32, min=1.0, max=1.0, mean=1.0),
  't': np.ndarray((200,), dtype=int64, min=0.0, max=34.0, mean=9.14)}

class ray.rllib.policy.sample_batch.SampleBatch(*args, **kwargs)
	def agent_steps() -> int  # 返回SampleBatch 中的 步长
	def env_steps() -> int #  和 agent_steps 相同
	def concat_samples() # 输入两个 SampleBatch， 可以将两个concat起来。
e.g. 
import numpy as np
from ray.rllib.policy.sample_batch import SampleBatch
b1 = SampleBatch({"a": np.array([1, 2]), "b": np.array([10, 11])})
b2 = SampleBatch({"a": np.array([3]),  "b": np.array([12])})
print(SampleBatch.concat_samples([b1, b2])) 

	def concat() # 返回一个新的SampleBatch
e.g.
import numpy as np
from ray.rllib.policy.sample_batch import SampleBatch
b1 = SampleBatch({"a": np.array([1, 2])}) 
b2 = SampleBatch({"a": np.array([3, 4, 5])}) 
print(b1.concat(b2)) 

	def copy(shallow: bool = False) → ray.rllib.policy.sample_batch.SampleBatch # 返回一个 浅拷贝/深拷贝的 SampleBatch

	def rows() # 生成一个迭代器，返回每步对应的pair
e.g.
from ray.rllib.policy.sample_batch import SampleBatch
batch = SampleBatch( {"a": [1, 2, 3], "b": [4, 5, 6],"seq_lens": [1, 2]} )
for row in batch.rows(): 
   print(row) 
>>{"a": 1, "b": 4, "seq_lens": 1}
>>{"a": 2, "b": 5, "seq_lens": 1}
>>{"a": 3, "b": 6, "seq_lens": 1}
>
	def shuffle() # 返回一个 被打乱 的 SampleBatch， 但是每个step的值都是对应的
	def split_by_episode() #  将sample结果按照episode if分开

Sampler

将不同worker生成的SampleBatch 进行合并收集
分为SyncSampler 和 AsyncSampler 。 在 WorkerSet中根据需要可以选择。RollWorker.sampler中就是这个类。

PolicyMap

用于多智能体中将policy id对应到成policy实例。

View requirement

RolloutWorker 在存储时 会按照step进行存储。 但是我们在想要将sample传入model中的时候，很多时候，一方面需要处理数据 ；另一方面可能会用到之前的数据。
如果需要处理数据，虽然rllib中有preprocessor，但是Ray建议 使用 wrapper对env进行包装，而不是创建自定义的preprocessor。
如果需要使用之前的数据，那么就需要用到 view requirement. 这是 sampler 和model 之间交流的设定方式。 view requirement 中定义 是否需要之前的数据。 如果需要，那么在model中计算动作compute_action之前，就会将需要的数据作为输入传入。

ViewRequirement class 包含四个变量：
data_col = “rewards”/"obs"等，用于标识之前数据中需要的数据。
space: gym.Space 说明数据的形式。
shift: int / list[int] 标识使用多会儿的数据。 0 标识这一时刻。 -1表示上一时刻。
used_for_training: bool 一般是True

如果需要，可以在model的 __init__中添加。