def Process(self, request_iterator, context):
"""Processes incoming EnvironmentRequests.
For each EnvironmentRequest the internal message is extracted and handled.
The response for that message is then placed in a EnvironmentResponse which
is returned to the client.
An error status will be returned if an unknown message type is received or
if the message is invalid for the current world state.
Args:
request_iterator: Message iterator provided by gRPC.
context: Context provided by gRPC.
Yields:
EnvironmentResponse: Response for each incoming EnvironmentRequest.
"""
env_factory = CatchGameFactory(_INITIAL_SEED)
env = None
is_joined = False
skip_next_frame = False
action_manager = spec_manager.SpecManager(_action_spec())
observation_manager = spec_manager.SpecManager(_observation_spec())
for request in request_iterator:
environment_response = dm_env_rpc_pb2.EnvironmentResponse()
try:
message_type = request.WhichOneof('payload')
internal_request = getattr(request, message_type)
_check_message_type(env, is_joined, message_type)
if message_type == 'create_world':
env = env_factory.new_game()
skip_next_frame = True
response = dm_env_rpc_pb2.CreateWorldResponse(
world_name=_WORLD_NAME)
elif message_type == 'join_world':
if is_joined:
raise RuntimeError(
f'Tried to join world "{internal_request.world_name}" but '
f'already joined to world "{_WORLD_NAME}"')
if internal_request.world_name != _WORLD_NAME:
raise RuntimeError(
f'Tried to join world "{internal_request.world_name}" but the '
f'only supported world is "{_WORLD_NAME}"')
response = dm_env_rpc_pb2.JoinWorldResponse()
for uid, action in _action_spec().items():
response.specs.actions[uid].CopyFrom(action)
for uid, observation in _observation_spec().items():
response.specs.observations[uid].CopyFrom(observation)
is_joined = True
elif message_type == 'step':
# We need to skip all actions after creating or resetting the
# environment.
if skip_next_frame:
skip_next_frame = False
else:
unpacked_actions = action_manager.unpack(
internal_request.actions)
paddle_action = unpacked_actions.get(
_ACTION_PADDLE, _DEFAULT_ACTION)
env.update(paddle_action)
response = dm_env_rpc_pb2.StepResponse()
packed_observations = observation_manager.pack({
_OBSERVATION_BOARD:
env.draw_board(),
_OBSERVATION_REWARD:
env.reward()
})
for requested_observation in internal_request.requested_observations:
response.observations[requested_observation].CopyFrom(
packed_observations[requested_observation])
if env.has_terminated():
response.state = dm_env_rpc_pb2.EnvironmentStateType.TERMINATED
else:
response.state = dm_env_rpc_pb2.EnvironmentStateType.RUNNING
if env.has_terminated():
env = env_factory.new_game()
skip_next_frame = True
elif message_type == 'reset':
env = env_factory.new_game()
skip_next_frame = True
response = dm_env_rpc_pb2.ResetResponse()
for uid, action in _action_spec().items():
response.specs.actions[uid].CopyFrom(action)
for uid, observation in _observation_spec().items():
response.specs.observations[uid].CopyFrom(observation)
elif message_type == 'reset_world':
env = env_factory.new_game()
skip_next_frame = True
response = dm_env_rpc_pb2.ResetWorldResponse()
elif message_type == 'leave_world':
is_joined = False
response = dm_env_rpc_pb2.LeaveWorldResponse()
elif message_type == 'destroy_world':
if internal_request.world_name != _WORLD_NAME:
raise RuntimeError(
'Tried to destroy world "{}" but we only support world "{}"'
.format(internal_request.world_name, _WORLD_NAME))
env = None
response = dm_env_rpc_pb2.DestroyWorldResponse()
else:
raise RuntimeError(
'Unhandled message: {}'.format(message_type))
getattr(environment_response, message_type).CopyFrom(response)
except Exception as e: # pylint: disable=broad-except
environment_response.error.CopyFrom(
status_pb2.Status(message=str(e)))
yield environment_response
def _pack_property_response(text_proto):
extension_message = any_pb2.Any()
extension_message.Pack(
text_format.Parse(text_proto, properties_pb2.PropertyResponse()))
return dm_env_rpc_pb2.EnvironmentResponse(extension=extension_message)
import mock
from google.protobuf import any_pb2
from google.protobuf import struct_pb2
from google.rpc import status_pb2
from dm_env_rpc.v1 import connection as dm_env_rpc_connection
from dm_env_rpc.v1 import dm_env_rpc_pb2
from dm_env_rpc.v1 import error
from dm_env_rpc.v1 import tensor_utils
_CREATE_REQUEST = dm_env_rpc_pb2.CreateWorldRequest(
settings={'foo': tensor_utils.pack_tensor('bar')})
_CREATE_RESPONSE = dm_env_rpc_pb2.CreateWorldResponse()
_BAD_CREATE_REQUEST = dm_env_rpc_pb2.CreateWorldRequest()
_TEST_ERROR = dm_env_rpc_pb2.EnvironmentResponse(error=status_pb2.Status(
message='A test error.'))
_INCORRECT_RESPONSE_TEST_MSG = dm_env_rpc_pb2.DestroyWorldRequest(
world_name='foo')
_INCORRECT_RESPONSE = dm_env_rpc_pb2.EnvironmentResponse(
leave_world=dm_env_rpc_pb2.LeaveWorldResponse())
_EXTENSION_REQUEST = struct_pb2.Value(string_value='extension request')
_EXTENSION_RESPONSE = struct_pb2.Value(number_value=555)
def _wrap_in_any(proto):
any_proto = any_pb2.Any()
any_proto.Pack(proto)
return any_proto
_pack_property_response("""list_property {
values: {
is_readable:true
spec { name: "baz.fiz" dtype:UINT32 shape: 2 shape: 2 }
}}"""),
_create_property_request_key('list_property {}'):
_pack_property_response("""list_property {
values: { is_readable:true spec { name: "foo" dtype:INT32 }
description: "This is a documented integer" }
values: { is_readable:true
is_writable:true
spec { name: "bar" dtype:STRING } }
values: { is_listable:true spec { name: "baz" } }
}"""),
_create_property_request_key('read_property { key: "bad_property" }'):
dm_env_rpc_pb2.EnvironmentResponse(error=status_pb2.Status(
message='invalid property request.'))
}
@contextlib.contextmanager
def _create_mock_connection():
"""Helper to create mock dm_env_rpc connection."""
with mock.patch.object(dm_env_rpc_connection,
'dm_env_rpc_pb2_grpc') as mock_grpc:
def _process(request_iterator):
for request in request_iterator:
yield _EXPECTED_REQUEST_RESPONSE_PAIRS[
request.SerializeToString()]
mock_stub_class = mock.MagicMock()