use async_trait::async_trait;
use rlmesh_grpc::env::{
CloseResponse as ProtoCloseResponse, Environment, RenderRequest as ProtoRenderRequest,
RenderResponse as ProtoRenderResponse, ResetRequest as ProtoResetRequest,
ResetResponse as ProtoResetResponse, StepRequest as ProtoStepRequest,
StepResponse as ProtoStepResponse,
};
use rlmesh_grpc::error::{EnvError, EnvErrorCode};
use rlmesh_grpc::wire::{
bytes_value, decode_batched_partial_values, encode_batched_partial_values,
meta_map_from_struct, meta_map_to_struct, render_result_to_proto, value_bytes,
};
use super::Env;
use super::types::{CloseRequest, EpisodeMetadata, RenderRequest, ResetRequest, StepRequest};
use crate::spaces;
pub struct WireEnvAdapter<E> {
inner: E,
}
impl<E> WireEnvAdapter<E> {
pub fn new(inner: E) -> Self {
Self { inner }
}
}
#[async_trait]
impl<E: Env> Environment for WireEnvAdapter<E> {
fn observation_space(&self) -> &spaces::SpaceSpec {
self.inner.observation_space()
}
fn action_space(&self) -> &spaces::SpaceSpec {
self.inner.action_space()
}
fn num_envs(&self) -> usize {
self.inner.num_envs()
}
fn env_contract(&self) -> &spaces::EnvContract {
self.inner.env_contract()
}
async fn reset(
&mut self,
req: ProtoResetRequest,
) -> std::result::Result<ProtoResetResponse, EnvError> {
let result = self
.inner
.reset(ResetRequest {
seeds: req.seeds,
options: req.options.map(meta_map_from_struct),
timeout_ms: req.timeout_ms,
})
.await
.map_err(gym_error_to_env_error)?;
validate_observation_count(&result.observations, self.inner.num_envs())?;
let observations =
encode_batched_partial_values(&result.observations, self.inner.observation_space())
.map_err(protocol_error_to_env_error)?;
Ok(ProtoResetResponse {
observation: Some(bytes_value(observations)),
infos: result.info.as_ref().map(meta_map_to_struct),
episode_ids: result.episode_ids,
})
}
async fn step(
&mut self,
req: ProtoStepRequest,
) -> std::result::Result<ProtoStepResponse, EnvError> {
let action_payload =
value_bytes(req.action.as_ref()).map_err(protocol_error_to_env_error)?;
let actions =
decode_batched_partial_values(action_payload.as_ref(), self.inner.action_space())
.map_err(protocol_error_to_env_error)?;
validate_action_count(&actions, self.inner.num_envs())?;
let result = self
.inner
.step(StepRequest {
actions,
timeout_ms: req.timeout_ms,
})
.await
.map_err(gym_error_to_env_error)?;
let env_count = self.inner.num_envs();
validate_observation_count(&result.observations, env_count)?;
validate_bool_count(&result.terminated, env_count, "terminated")?;
validate_bool_count(&result.truncated, env_count, "truncated")?;
validate_f64_count(&result.rewards, env_count, "rewards")?;
let observations =
encode_batched_partial_values(&result.observations, self.inner.observation_space())
.map_err(protocol_error_to_env_error)?;
Ok(ProtoStepResponse {
observation: Some(bytes_value(observations)),
rewards: result.rewards,
terminated_mask: result.terminated.into_iter().map(u8::from).collect(),
truncated_mask: result.truncated.into_iter().map(u8::from).collect(),
infos: result.info.as_ref().map(meta_map_to_struct),
completed_episodes: result
.completed_episodes
.iter()
.map(public_episode_metadata_to_proto)
.collect::<std::result::Result<Vec<_>, _>>()?,
episode_ids: result.episode_ids,
})
}
async fn render(
&mut self,
req: ProtoRenderRequest,
) -> std::result::Result<ProtoRenderResponse, EnvError> {
let result = self
.inner
.render(RenderRequest {
env_index: render_env_index(&req.mask)?,
timeout_ms: req.timeout_ms,
})
.await
.map_err(gym_error_to_env_error)?;
Ok(render_result_to_proto(&result))
}
async fn close(&mut self) -> std::result::Result<ProtoCloseResponse, EnvError> {
let result = self
.inner
.close(CloseRequest {
wait_for_episodes: false,
})
.await
.map_err(gym_error_to_env_error)?;
Ok(ProtoCloseResponse {
final_episodes: result
.final_episodes
.iter()
.map(public_episode_metadata_to_proto)
.collect::<std::result::Result<Vec<_>, _>>()?,
})
}
}
fn gym_error_to_env_error(error: spaces::EnvRuntimeError) -> EnvError {
match error {
spaces::EnvRuntimeError::InvalidSpace(message)
| spaces::EnvRuntimeError::InvalidValue(message) => {
EnvError::new(EnvErrorCode::InvalidAction, message)
}
spaces::EnvRuntimeError::Runtime(message) => EnvError::new(EnvErrorCode::Internal, message),
}
}
fn public_episode_metadata_to_proto(
value: &EpisodeMetadata,
) -> std::result::Result<rlmesh_proto::env::v1::EpisodeMetadata, EnvError> {
Ok(rlmesh_proto::env::v1::EpisodeMetadata {
episode_id: value.episode_id.clone(),
seed: value.seed,
env_index: value.env_index,
step_count: value.step_count,
cumulative_reward: value.cumulative_reward,
terminated: value.terminated,
truncated: value.truncated,
start_timestamp_ns: value.start_timestamp_ns,
end_timestamp_ns: value.end_timestamp_ns,
duration_ms: value.duration_ms,
final_info: value.final_info.as_ref().map(meta_map_to_struct),
})
}
pub(super) fn protocol_error_to_error(error: impl ToString) -> crate::Error {
crate::Error::Internal(error.to_string())
}
pub(super) fn proto_episode_metadata_to_public(
value: rlmesh_proto::env::v1::EpisodeMetadata,
) -> std::result::Result<EpisodeMetadata, String> {
Ok(EpisodeMetadata {
episode_id: value.episode_id,
seed: value.seed,
env_index: value.env_index,
step_count: value.step_count,
cumulative_reward: value.cumulative_reward,
terminated: value.terminated,
truncated: value.truncated,
start_timestamp_ns: value.start_timestamp_ns,
end_timestamp_ns: value.end_timestamp_ns,
duration_ms: value.duration_ms,
final_info: value.final_info.map(meta_map_from_struct),
})
}
fn render_env_index(mask: &[u8]) -> std::result::Result<Option<usize>, EnvError> {
let indices = mask
.iter()
.enumerate()
.filter_map(|(index, value)| (*value != 0).then_some(index))
.collect::<Vec<_>>();
match indices.as_slice() {
[] => Ok(None),
[index] => Ok(Some(*index)),
_ => Err(EnvError::new(
EnvErrorCode::InvalidAction,
"render requests support at most one env_index".to_string(),
)),
}
}
pub(super) fn validate_action_count(
actions: &[spaces::SpaceValue],
num_envs: usize,
) -> std::result::Result<(), EnvError> {
if actions.len() == num_envs {
return Ok(());
}
Err(EnvError::new(
EnvErrorCode::InvalidAction,
format!("expected {num_envs} actions, got {}", actions.len()),
))
}
pub(super) fn validate_observation_count(
observations: &[spaces::SpaceValue],
num_envs: usize,
) -> std::result::Result<(), EnvError> {
if observations.len() == num_envs {
return Ok(());
}
Err(EnvError::new(
EnvErrorCode::Internal,
format!(
"env returned {} observations for {num_envs} envs",
observations.len()
),
))
}
pub(super) fn validate_bool_count(
values: &[bool],
num_envs: usize,
label: &str,
) -> std::result::Result<(), EnvError> {
if values.len() == num_envs {
return Ok(());
}
Err(EnvError::new(
EnvErrorCode::Internal,
format!(
"env returned {} {label} values for {num_envs} envs",
values.len()
),
))
}
pub(super) fn validate_f64_count(
values: &[f64],
num_envs: usize,
label: &str,
) -> std::result::Result<(), EnvError> {
if values.len() == num_envs {
return Ok(());
}
Err(EnvError::new(
EnvErrorCode::Internal,
format!(
"env returned {} {label} values for {num_envs} envs",
values.len()
),
))
}
fn protocol_error_to_env_error(error: impl ToString) -> EnvError {
EnvError::new(EnvErrorCode::Internal, error.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use crate::{BindAddress, ServeOptions};
use super::super::{
CloseResult, EnvServer, RemoteEnv, RenderResult, ResetRequest, ResetResult, StepResult,
};
struct DummyEnv {
obs_space: spaces::SpaceSpec,
action_space: spaces::SpaceSpec,
env_contract: spaces::EnvContract,
last_render_request: Option<RenderRequest>,
closes: Option<Arc<AtomicUsize>>,
}
impl DummyEnv {
fn new() -> Self {
Self::new_with_close_counter(None)
}
fn new_with_close_counter(closes: Option<Arc<AtomicUsize>>) -> Self {
let obs_space = spaces::spaces::BoxSpaceBuilder::scalar(-1.0, 1.0, vec![2])
.dtype(spaces::DType::Float32)
.build()
.unwrap();
let action_space = spaces::spaces::DiscreteBuilder::new(3).build().unwrap();
let env_contract = spaces::EnvContract {
id: "DummyEnv-v1".to_string(),
observation_space: Some(obs_space.clone()),
action_space: Some(action_space.clone()),
metadata: None,
render_mode: String::new(),
num_envs: 2,
};
Self {
obs_space,
action_space,
env_contract,
last_render_request: None,
closes,
}
}
}
#[async_trait]
impl Env for DummyEnv {
fn observation_space(&self) -> &spaces::SpaceSpec {
&self.obs_space
}
fn action_space(&self) -> &spaces::SpaceSpec {
&self.action_space
}
fn num_envs(&self) -> usize {
2
}
fn env_contract(&self) -> &spaces::EnvContract {
&self.env_contract
}
async fn reset(
&mut self,
_req: ResetRequest,
) -> std::result::Result<ResetResult, spaces::EnvRuntimeError> {
Ok(ResetResult {
observations: vec![
spaces::SpaceValue::Box(spaces::BoxValue::new(
vec![0; 8],
vec![2],
spaces::DType::Float32,
)),
spaces::SpaceValue::Box(spaces::BoxValue::new(
vec![1; 8],
vec![2],
spaces::DType::Float32,
)),
],
info: None,
episode_ids: vec!["ep-0".to_string(), "ep-1".to_string()],
})
}
async fn step(
&mut self,
req: StepRequest,
) -> std::result::Result<StepResult, spaces::EnvRuntimeError> {
Ok(StepResult {
observations: req
.actions
.into_iter()
.map(|action| match action {
spaces::SpaceValue::Discrete(value) => {
spaces::SpaceValue::Box(spaces::BoxValue::new(
vec![value as u8; 8],
vec![2],
spaces::DType::Float32,
))
}
other => other,
})
.collect(),
rewards: vec![1.0, 2.0],
terminated: vec![false, true],
truncated: vec![false, false],
info: None,
completed_episodes: vec![],
episode_ids: vec!["ep-0".to_string(), "ep-1b".to_string()],
})
}
async fn render(
&mut self,
req: RenderRequest,
) -> std::result::Result<RenderResult, spaces::EnvRuntimeError> {
self.last_render_request = Some(req);
Ok(RenderResult {
frame: Some(spaces::RenderFrame {
png_frame: vec![1, 2, 3],
}),
})
}
async fn close(
&mut self,
_req: CloseRequest,
) -> std::result::Result<CloseResult, spaces::EnvRuntimeError> {
if let Some(closes) = &self.closes {
closes.fetch_add(1, Ordering::SeqCst);
}
Ok(CloseResult {
final_episodes: vec![],
})
}
}
#[tokio::test]
async fn wire_adapter_roundtrips_batched_reset_and_step() {
let mut env = WireEnvAdapter::new(DummyEnv::new());
let reset = Environment::reset(
&mut env,
ProtoResetRequest {
seeds: vec![7, 8],
..Default::default()
},
)
.await
.unwrap();
let reset_payload = value_bytes(reset.observation.as_ref()).unwrap().unwrap();
let reset_obs =
decode_batched_partial_values(Some(&reset_payload), env.observation_space()).unwrap();
assert_eq!(reset_obs.len(), 2);
let actions = [
spaces::SpaceValue::Discrete(1),
spaces::SpaceValue::Discrete(2),
];
let action_space = env.action_space().clone();
let step = Environment::step(
&mut env,
ProtoStepRequest {
action: Some(bytes_value(
encode_batched_partial_values(&actions, &action_space).unwrap(),
)),
timeout_ms: 0,
},
)
.await
.unwrap();
let step_payload = value_bytes(step.observation.as_ref()).unwrap().unwrap();
let step_obs =
decode_batched_partial_values(Some(&step_payload), env.observation_space()).unwrap();
assert_eq!(step_obs.len(), 2);
assert_eq!(step.rewards, vec![1.0, 2.0]);
assert_eq!(step.terminated_mask, vec![0, 1]);
}
#[tokio::test]
async fn wire_adapter_rejects_wrong_action_count() {
let mut env = WireEnvAdapter::new(DummyEnv::new());
let actions = [spaces::SpaceValue::Discrete(1)];
let action_space = env.action_space().clone();
let error = Environment::step(
&mut env,
ProtoStepRequest {
action: Some(bytes_value(
encode_batched_partial_values(&actions, &action_space).unwrap(),
)),
timeout_ms: 0,
},
)
.await
.unwrap_err();
assert_eq!(error.code, EnvErrorCode::InvalidAction);
}
#[tokio::test]
async fn wire_adapter_maps_render_mask_to_env_index() {
let mut env = WireEnvAdapter::new(DummyEnv::new());
let result = Environment::render(
&mut env,
ProtoRenderRequest {
mask: vec![0, 1],
timeout_ms: 0,
},
)
.await
.unwrap();
assert!(result.png_frame.is_some());
}
#[tokio::test]
async fn served_env_close_detaches_and_shutdown_stops_server() {
let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
.await
.unwrap();
let port = listener.local_addr().unwrap().port();
drop(listener);
let closes = Arc::new(AtomicUsize::new(0));
let server_closes = Arc::clone(&closes);
let server = tokio::spawn(async move {
EnvServer::new(DummyEnv::new_with_close_counter(Some(server_closes)))
.serve_with_options(
BindAddress::Tcp {
host: "127.0.0.1".to_string(),
port,
},
ServeOptions {
allow_remote_shutdown: true,
..ServeOptions::default()
},
)
.await
});
let address = format!("tcp://127.0.0.1:{port}");
let mut client = loop {
match RemoteEnv::connect(&address).await {
Ok(client) => break client,
Err(err) if !server.is_finished() => {
let _ = err;
tokio::time::sleep(Duration::from_millis(10)).await;
}
Err(err) => panic!("environment server did not start: {err}"),
}
};
client.close().await.unwrap();
assert_eq!(closes.load(Ordering::SeqCst), 0);
let mut second_client = RemoteEnv::connect(&address).await.unwrap();
let _ = second_client.reset(ResetRequest::default()).await.unwrap();
assert!(second_client.shutdown("test shutdown").await.unwrap());
tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap()
.unwrap();
assert_eq!(closes.load(Ordering::SeqCst), 1);
}
}