from google.protobuf import struct_pb2 as _struct_pb2
from . import gym_pb2 as _gym_pb2
from google.protobuf.internal import containers as _containers
from google.protobuf import descriptor as _descriptor
from google.protobuf import message as _message
from collections.abc import Iterable as _Iterable, Mapping as _Mapping
from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union

DESCRIPTOR: _descriptor.FileDescriptor

class Tensor(_message.Message):
    __slots__ = ("dtype", "shape", "raw")
    DTYPE_FIELD_NUMBER: _ClassVar[int]
    SHAPE_FIELD_NUMBER: _ClassVar[int]
    RAW_FIELD_NUMBER: _ClassVar[int]
    dtype: _gym_pb2.DType
    shape: _containers.RepeatedScalarFieldContainer[int]
    raw: bytes
    def __init__(self, dtype: _Optional[_Union[_gym_pb2.DType, str]] = ..., shape: _Optional[_Iterable[int]] = ..., raw: _Optional[bytes] = ...) -> None: ...

class ResetReq(_message.Message):
    __slots__ = ("seeds",)
    SEEDS_FIELD_NUMBER: _ClassVar[int]
    seeds: _containers.RepeatedScalarFieldContainer[int]
    def __init__(self, seeds: _Optional[_Iterable[int]] = ...) -> None: ...

class StepReq(_message.Message):
    __slots__ = ("actions",)
    ACTIONS_FIELD_NUMBER: _ClassVar[int]
    actions: Tensor
    def __init__(self, actions: _Optional[_Union[Tensor, _Mapping]] = ...) -> None: ...

class RenderReq(_message.Message):
    __slots__ = ()
    def __init__(self) -> None: ...

class CloseReq(_message.Message):
    __slots__ = ()
    def __init__(self) -> None: ...

class StateReq(_message.Message):
    __slots__ = ()
    def __init__(self) -> None: ...

class ClientMsg(_message.Message):
    __slots__ = ("reset", "step", "render", "close", "state")
    RESET_FIELD_NUMBER: _ClassVar[int]
    STEP_FIELD_NUMBER: _ClassVar[int]
    RENDER_FIELD_NUMBER: _ClassVar[int]
    CLOSE_FIELD_NUMBER: _ClassVar[int]
    STATE_FIELD_NUMBER: _ClassVar[int]
    reset: ResetReq
    step: StepReq
    render: RenderReq
    close: CloseReq
    state: StateReq
    def __init__(self, reset: _Optional[_Union[ResetReq, _Mapping]] = ..., step: _Optional[_Union[StepReq, _Mapping]] = ..., render: _Optional[_Union[RenderReq, _Mapping]] = ..., close: _Optional[_Union[CloseReq, _Mapping]] = ..., state: _Optional[_Union[StateReq, _Mapping]] = ...) -> None: ...

class ResetOk(_message.Message):
    __slots__ = ("obs", "infos")
    OBS_FIELD_NUMBER: _ClassVar[int]
    INFOS_FIELD_NUMBER: _ClassVar[int]
    obs: Tensor
    infos: _struct_pb2.Struct
    def __init__(self, obs: _Optional[_Union[Tensor, _Mapping]] = ..., infos: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...

class StepOk(_message.Message):
    __slots__ = ("obs", "rewards", "terminated_mask", "truncated_mask", "infos")
    OBS_FIELD_NUMBER: _ClassVar[int]
    REWARDS_FIELD_NUMBER: _ClassVar[int]
    TERMINATED_MASK_FIELD_NUMBER: _ClassVar[int]
    TRUNCATED_MASK_FIELD_NUMBER: _ClassVar[int]
    INFOS_FIELD_NUMBER: _ClassVar[int]
    obs: Tensor
    rewards: Tensor
    terminated_mask: bytes
    truncated_mask: bytes
    infos: _struct_pb2.Struct
    def __init__(self, obs: _Optional[_Union[Tensor, _Mapping]] = ..., rewards: _Optional[_Union[Tensor, _Mapping]] = ..., terminated_mask: _Optional[bytes] = ..., truncated_mask: _Optional[bytes] = ..., infos: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...

class RenderOk(_message.Message):
    __slots__ = ("frame",)
    FRAME_FIELD_NUMBER: _ClassVar[int]
    frame: Tensor
    def __init__(self, frame: _Optional[_Union[Tensor, _Mapping]] = ...) -> None: ...

class EpisodeResult(_message.Message):
    __slots__ = ("score", "duration", "video_key")
    SCORE_FIELD_NUMBER: _ClassVar[int]
    DURATION_FIELD_NUMBER: _ClassVar[int]
    VIDEO_KEY_FIELD_NUMBER: _ClassVar[int]
    score: float
    duration: float
    video_key: str
    def __init__(self, score: _Optional[float] = ..., duration: _Optional[float] = ..., video_key: _Optional[str] = ...) -> None: ...

class CloseOk(_message.Message):
    __slots__ = ("score", "duration", "episode_results")
    SCORE_FIELD_NUMBER: _ClassVar[int]
    DURATION_FIELD_NUMBER: _ClassVar[int]
    EPISODE_RESULTS_FIELD_NUMBER: _ClassVar[int]
    score: float
    duration: float
    episode_results: _containers.RepeatedCompositeFieldContainer[EpisodeResult]
    def __init__(self, score: _Optional[float] = ..., duration: _Optional[float] = ..., episode_results: _Optional[_Iterable[_Union[EpisodeResult, _Mapping]]] = ...) -> None: ...

class ServerError(_message.Message):
    __slots__ = ("message",)
    MESSAGE_FIELD_NUMBER: _ClassVar[int]
    message: str
    def __init__(self, message: _Optional[str] = ...) -> None: ...

class ServerMsg(_message.Message):
    __slots__ = ("reset_ok", "step_ok", "render_ok", "close_ok", "error")
    RESET_OK_FIELD_NUMBER: _ClassVar[int]
    STEP_OK_FIELD_NUMBER: _ClassVar[int]
    RENDER_OK_FIELD_NUMBER: _ClassVar[int]
    CLOSE_OK_FIELD_NUMBER: _ClassVar[int]
    ERROR_FIELD_NUMBER: _ClassVar[int]
    reset_ok: ResetOk
    step_ok: StepOk
    render_ok: RenderOk
    close_ok: CloseOk
    error: ServerError
    def __init__(self, reset_ok: _Optional[_Union[ResetOk, _Mapping]] = ..., step_ok: _Optional[_Union[StepOk, _Mapping]] = ..., render_ok: _Optional[_Union[RenderOk, _Mapping]] = ..., close_ok: _Optional[_Union[CloseOk, _Mapping]] = ..., error: _Optional[_Union[ServerError, _Mapping]] = ...) -> None: ...
