wehub-resource-sync
127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
import pickle
|
|
import socket
|
|
import time
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
|
|
from ray.rllib.core import (
|
|
COMPONENT_RL_MODULE,
|
|
Columns,
|
|
)
|
|
from ray.rllib.env.external.rllink import (
|
|
RLlink,
|
|
get_rllink_message,
|
|
send_rllink_message,
|
|
)
|
|
from ray.rllib.env.single_agent_episode import SingleAgentEpisode
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
from ray.rllib.utils.numpy import softmax
|
|
|
|
torch, _ = try_import_torch()
|
|
|
|
|
|
def _dummy_external_client(port: int = 5556):
|
|
"""A dummy client that runs CartPole and acts as a testing external env."""
|
|
|
|
def _set_state(msg_body, rl_module):
|
|
rl_module.set_state(msg_body[COMPONENT_RL_MODULE])
|
|
# return msg_body[WEIGHTS_SEQ_NO]
|
|
|
|
# Connect to server.
|
|
while True:
|
|
try:
|
|
print(f"Trying to connect to localhost:{port} ...")
|
|
sock_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock_.connect(("localhost", port))
|
|
break
|
|
except ConnectionRefusedError:
|
|
time.sleep(5)
|
|
|
|
# Send ping-pong.
|
|
send_rllink_message(sock_, {"type": RLlink.PING.name})
|
|
msg_type, msg_body = get_rllink_message(sock_)
|
|
assert msg_type == RLlink.PONG
|
|
|
|
# Request config.
|
|
send_rllink_message(sock_, {"type": RLlink.GET_CONFIG.name})
|
|
msg_type, msg_body = get_rllink_message(sock_)
|
|
assert msg_type == RLlink.SET_CONFIG
|
|
|
|
config = pickle.loads(msg_body["config"])
|
|
# Create the RLModule.
|
|
rl_module = config.get_rl_module_spec().build()
|
|
|
|
# Request state/weights.
|
|
send_rllink_message(sock_, {"type": RLlink.GET_STATE.name})
|
|
msg_type, msg_body = get_rllink_message(sock_)
|
|
assert msg_type == RLlink.SET_STATE
|
|
_set_state(msg_body["state"], rl_module)
|
|
|
|
env_steps_per_sample = config.get_rollout_fragment_length()
|
|
|
|
# Start actual env loop.
|
|
env = gym.make("CartPole-v1")
|
|
obs, _ = env.reset()
|
|
episode = SingleAgentEpisode(observations=[obs])
|
|
episodes = [episode]
|
|
|
|
while True:
|
|
# Perform action inference using the RLModule.
|
|
logits = rl_module.forward_exploration(
|
|
batch={
|
|
Columns.OBS: torch.tensor(np.array([obs], np.float32)),
|
|
}
|
|
)[Columns.ACTION_DIST_INPUTS][
|
|
0
|
|
].numpy() # [0]=batch size 1
|
|
|
|
# Stochastic sample.
|
|
action_probs = softmax(logits)
|
|
action = int(np.random.choice(list(range(env.action_space.n)), p=action_probs))
|
|
logp = float(np.log(action_probs[action]))
|
|
|
|
# Perform the env step.
|
|
obs, reward, terminated, truncated, _ = env.step(action)
|
|
|
|
# Collect step data.
|
|
episode.add_env_step(
|
|
action=action,
|
|
reward=reward,
|
|
observation=obs,
|
|
terminated=terminated,
|
|
truncated=truncated,
|
|
extra_model_outputs={
|
|
Columns.ACTION_DIST_INPUTS: logits,
|
|
Columns.ACTION_LOGP: logp,
|
|
},
|
|
)
|
|
|
|
# We collected enough samples -> Send them to server.
|
|
if sum(map(len, episodes)) == env_steps_per_sample:
|
|
# Send the data to the server.
|
|
send_rllink_message(
|
|
sock_,
|
|
{
|
|
"type": RLlink.EPISODES_AND_GET_STATE.name,
|
|
"episodes": [e.get_state() for e in episodes],
|
|
"timesteps": env_steps_per_sample,
|
|
},
|
|
)
|
|
# We are forced to sample on-policy. Have to wait for a response
|
|
# with the state (weights) in it.
|
|
msg_type, msg_body = get_rllink_message(sock_)
|
|
assert msg_type == RLlink.SET_STATE
|
|
_set_state(msg_body["state"], rl_module)
|
|
|
|
episodes = []
|
|
if not episode.is_done:
|
|
episode = episode.cut()
|
|
episodes.append(episode)
|
|
|
|
# If episode is done, reset env and create a new episode.
|
|
if episode.is_done:
|
|
obs, _ = env.reset()
|
|
episode = SingleAgentEpisode(observations=[obs])
|
|
episodes.append(episode)
|