"""REST client to interact with a policy server.
This client supports both local and remote policy inference modes. Local
inference is faster but causes more compute to be done on the client.
"""
import logging
import threading
import time
from typing import Union, Optional
from enum import Enum
import ray.cloudpickle as pickle
from ray.rllib.env.external_env import ExternalEnv
from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv
from ray.rllib.env.multi_agent_env import MultiAgentEnv
from ray.rllib.policy.sample_batch import MultiAgentBatch
from ray.rllib.utils.annotations import PublicAPI
from ray.rllib.utils.typing import (
MultiAgentDict,
EnvInfoDict,
EnvObsType,
EnvActionType,
)
logger = logging.getLogger(__name__)
try:
import requests # `requests` is not part of stdlib.
except ImportError:
requests = None
logger.warning(
"Couldn't import `requests` library. Be sure to install it on"
" the client side."
)
@PublicAPI
class Commands(Enum):
# Generic commands (for both modes).
ACTION_SPACE = "ACTION_SPACE"
OBSERVATION_SPACE = "OBSERVATION_SPACE"
# Commands for local inference mode.
GET_WORKER_ARGS = "GET_WORKER_ARGS"
GET_WEIGHTS = "GET_WEIGHTS"
REPORT_SAMPLES = "REPORT_SAMPLES"
# Commands for remote inference mode.
START_EPISODE = "START_EPISODE"
GET_ACTION = "GET_ACTION"
LOG_ACTION = "LOG_ACTION"
LOG_RETURNS = "LOG_RETURNS"
END_EPISODE = "END_EPISODE"
[docs]
@PublicAPI
class PolicyClient:
"""REST client to interact with an RLlib policy server."""
@PublicAPI
def __init__(
self,
address: str,
inference_mode: str = "local",
update_interval: float = 10.0,
session: Optional[requests.Session] = None,
):
"""Create a PolicyClient instance.
Args:
address: Server to connect to (e.g., "localhost:9090").
inference_mode: Whether to use 'local' or 'remote' policy
inference for computing actions.
update_interval (float or None): If using 'local' inference mode,
the policy is refreshed after this many seconds have passed,
or None for manual control via client.
session (requests.Session or None): If available the session object
is used to communicate with the policy server. Using a session
can lead to speedups as connections are reused. It is the
responsibility of the creator of the session to close it.
"""
self.address = address
self.session = session
self.env: ExternalEnv = None
if inference_mode == "local":
self.local = True
self._setup_local_rollout_worker(update_interval)
elif inference_mode == "remote":
self.local = False
else:
raise ValueError("inference_mode must be either 'local' or 'remote'")
[docs]
@PublicAPI
def start_episode(
self, episode_id: Optional[str] = None, training_enabled: bool = True
) -> str:
"""Record the start of one or more episode(s).
Args:
episode_id (Optional[str]): Unique string id for the episode or
None for it to be auto-assigned.
training_enabled: Whether to use experiences for this
episode to improve the policy.
Returns:
episode_id: Unique string id for the episode.
"""
if self.local:
self._update_local_policy()
return self.env.start_episode(episode_id, training_enabled)
return self._send(
{
"episode_id": episode_id,
"command": Commands.START_EPISODE,
"training_enabled": training_enabled,
}
)["episode_id"]
[docs]
@PublicAPI
def get_action(
self, episode_id: str, observation: Union[EnvObsType, MultiAgentDict]
) -> Union[EnvActionType, MultiAgentDict]:
"""Record an observation and get the on-policy action.
Args:
episode_id: Episode id returned from start_episode().
observation: Current environment observation.
Returns:
action: Action from the env action space.
"""
if self.local:
self._update_local_policy()
if isinstance(episode_id, (list, tuple)):
actions = {
eid: self.env.get_action(eid, observation[eid])
for eid in episode_id
}
return actions
else:
return self.env.get_action(episode_id, observation)
else:
return self._send(
{
"command": Commands.GET_ACTION,
"observation": observation,
"episode_id": episode_id,
}
)["action"]
[docs]
@PublicAPI
def log_action(
self,
episode_id: str,
observation: Union[EnvObsType, MultiAgentDict],
action: Union[EnvActionType, MultiAgentDict],
) -> None:
"""Record an observation and (off-policy) action taken.
Args:
episode_id: Episode id returned from start_episode().
observation: Current environment observation.
action: Action for the observation.
"""
if self.local:
self._update_local_policy()
return self.env.log_action(episode_id, observation, action)
self._send(
{
"command": Commands.LOG_ACTION,
"observation": observation,
"action": action,
"episode_id": episode_id,
}
)
[docs]
@PublicAPI
def log_returns(
self,
episode_id: str,
reward: float,
info: Union[EnvInfoDict, MultiAgentDict] = None,
multiagent_done_dict: Optional[MultiAgentDict] = None,
) -> None:
"""Record returns from the environment.
The reward will be attributed to the previous action taken by the
episode. Rewards accumulate until the next action. If no reward is
logged before the next action, a reward of 0.0 is assumed.
Args:
episode_id: Episode id returned from start_episode().
reward: Reward from the environment.
info: Extra info dict.
multiagent_done_dict: Multi-agent done information.
"""
if self.local:
self._update_local_policy()
if multiagent_done_dict is not None:
assert isinstance(reward, dict)
return self.env.log_returns(
episode_id, reward, info, multiagent_done_dict
)
return self.env.log_returns(episode_id, reward, info)
self._send(
{
"command": Commands.LOG_RETURNS,
"reward": reward,
"info": info,
"episode_id": episode_id,
"done": multiagent_done_dict,
}
)
[docs]
@PublicAPI
def end_episode(
self, episode_id: str, observation: Union[EnvObsType, MultiAgentDict]
) -> None:
"""Record the end of an episode.
Args:
episode_id: Episode id returned from start_episode().
observation: Current environment observation.
"""
if self.local:
self._update_local_policy()
return self.env.end_episode(episode_id, observation)
self._send(
{
"command": Commands.END_EPISODE,
"observation": observation,
"episode_id": episode_id,
}
)
[docs]
@PublicAPI
def update_policy_weights(self) -> None:
"""Query the server for new policy weights, if local inference is enabled."""
self._update_local_policy(force=True)
def _send(self, data):
payload = pickle.dumps(data)
if self.session is None:
response = requests.post(self.address, data=payload)
else:
response = self.session.post(self.address, data=payload)
if response.status_code != 200:
logger.error("Request failed {}: {}".format(response.text, data))
response.raise_for_status()
parsed = pickle.loads(response.content)
return parsed
def _setup_local_rollout_worker(self, update_interval):
self.update_interval = update_interval
self.last_updated = 0
logger.info("Querying server for rollout worker settings.")
kwargs = self._send(
{
"command": Commands.GET_WORKER_ARGS,
}
)["worker_args"]
(self.rollout_worker, self.inference_thread) = _create_embedded_rollout_worker(
kwargs, self._send
)
self.env = self.rollout_worker.env
def _update_local_policy(self, force=False):
assert self.inference_thread.is_alive()
if (
self.update_interval
and time.time() - self.last_updated > self.update_interval
) or force:
logger.info("Querying server for new policy weights.")
resp = self._send(
{
"command": Commands.GET_WEIGHTS,
}
)
weights = resp["weights"]
global_vars = resp["global_vars"]
logger.info(
"Updating rollout worker weights and global vars {}.".format(
global_vars
)
)
self.rollout_worker.set_weights(weights, global_vars)
self.last_updated = time.time()
class _LocalInferenceThread(threading.Thread):
"""Thread that handles experience generation (worker.sample() loop)."""
def __init__(self, rollout_worker, send_fn):
super().__init__()
self.daemon = True
self.rollout_worker = rollout_worker
self.send_fn = send_fn
def run(self):
try:
while True:
logger.info("Generating new batch of experiences.")
samples = self.rollout_worker.sample()
metrics = self.rollout_worker.get_metrics()
if isinstance(samples, MultiAgentBatch):
logger.info(
"Sending batch of {} env steps ({} agent steps) to "
"server.".format(samples.env_steps(), samples.agent_steps())
)
else:
logger.info(
"Sending batch of {} steps back to server.".format(
samples.count
)
)
self.send_fn(
{
"command": Commands.REPORT_SAMPLES,
"samples": samples,
"metrics": metrics,
}
)
except Exception as e:
logger.error("Error: inference worker thread died!", e)
def _auto_wrap_external(real_env_creator):
"""Wrap an environment in the ExternalEnv interface if needed.
Args:
real_env_creator: Create an env given the env_config.
"""
def wrapped_creator(env_config):
real_env = real_env_creator(env_config)
if not isinstance(real_env, (ExternalEnv, ExternalMultiAgentEnv)):
logger.info(
"The env you specified is not a supported (sub-)type of "
"ExternalEnv. Attempting to convert it automatically to "
"ExternalEnv."
)
if isinstance(real_env, MultiAgentEnv):
external_cls = ExternalMultiAgentEnv
else:
external_cls = ExternalEnv
class _ExternalEnvWrapper(external_cls):
def __init__(self, real_env):
super().__init__(
observation_space=real_env.observation_space,
action_space=real_env.action_space,
)
def run(self):
# Since we are calling methods on this class in the
# client, run doesn't need to do anything.
time.sleep(999999)
return _ExternalEnvWrapper(real_env)
return real_env
return wrapped_creator
def _create_embedded_rollout_worker(kwargs, send_fn):
"""Create a local rollout worker and a thread that samples from it.
Args:
kwargs: Args for the RolloutWorker constructor.
send_fn: Function to send a JSON request to the server.
"""
# Since the server acts as an input datasource, we have to reset the
# input config to the default, which runs env rollouts.
kwargs = kwargs.copy()
kwargs["config"] = kwargs["config"].copy(copy_frozen=False)
config = kwargs["config"]
config.output = None
config.input_ = "sampler"
config.input_config = {}
# If server has no env (which is the expected case):
# Generate a dummy ExternalEnv here using RandomEnv and the
# given observation/action spaces.
if config.env is None:
from ray.rllib.examples.envs.classes.random_env import (
RandomEnv,
RandomMultiAgentEnv,
)
env_config = {
"action_space": config.action_space,
"observation_space": config.observation_space,
}
is_ma = config.is_multi_agent()
kwargs["env_creator"] = _auto_wrap_external(
lambda _: (RandomMultiAgentEnv if is_ma else RandomEnv)(env_config)
)
# kwargs["config"].env = True
# Otherwise, use the env specified by the server args.
else:
real_env_creator = kwargs["env_creator"]
kwargs["env_creator"] = _auto_wrap_external(real_env_creator)
logger.info("Creating rollout worker with kwargs={}".format(kwargs))
from ray.rllib.evaluation.rollout_worker import RolloutWorker
rollout_worker = RolloutWorker(**kwargs)
inference_thread = _LocalInferenceThread(rollout_worker, send_fn)
inference_thread.start()
return rollout_worker, inference_thread