relayrl_framework 0.5.0-beta.2

A system-oriented multi-agent reinforcement learning framework.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

use crate::network::client::runtime::router::buffer::TrajectorySinkError;
use crate::network::client::runtime::router::filter::FilterError;
#[cfg(any(feature = "nats-transport", feature = "zmq-transport"))]
use crate::network::client::runtime::router::receiver::TransportReceiverError;

use relayrl_types::data::action::RelayRLAction;
use relayrl_types::data::trajectory::RelayRLTrajectory;

use active_uuid_registry::registry_uuid::Uuid;

use std::any::Any;
use std::sync::Arc;
use thiserror::Error;
use tokio::sync::oneshot;

pub(crate) mod buffer;
pub(crate) mod filter;
#[cfg(any(feature = "nats-transport", feature = "zmq-transport"))]
pub(crate) mod receiver;
pub(crate) mod router_dispatcher;

#[derive(Debug, Error)]
#[allow(clippy::enum_variant_names)]
pub enum RouterError {
    #[error(transparent)]
    FilterError(#[from] FilterError),
    #[error(transparent)]
    #[cfg(any(feature = "nats-transport", feature = "zmq-transport"))]
    TransportReceiverError(#[from] TransportReceiverError),
    #[error(transparent)]
    TrajectorySinkError(#[from] TrajectorySinkError),
}

pub(crate) struct RoutedMessage {
    pub actor_id: Uuid,
    pub protocol: RoutingProtocol,
    pub payload: RoutedPayload,
}

pub(crate) enum RoutingProtocol {
    ModelHandshake,
    RequestInference,
    RequestInferenceBatch,
    FlagLastInference,
    ModelVersion,
    ModelUpdate,
    SendTrajectory,
    Shutdown,
}

pub(crate) enum RoutedPayload {
    ModelHandshake,
    RequestInference(Box<InferenceRequest>),
    RequestInferenceBatch(Box<BatchedInferenceRequest>),
    FlagLastInference {
        reward: f32,
        env_id: Option<Uuid>,
        env_label: Option<String>,
    },
    ModelVersion {
        reply_to: oneshot::Sender<i64>,
    },
    ModelUpdate {
        model_bytes: Vec<u8>,
        version: i64,
    },
    SendTrajectory {
        timestamp: (u128, u128),
        trajectory: RelayRLTrajectory,
    },
    Shutdown,
}

/// observation and mask are Arc<AnyBurnTensor<B, D_IN>> and Arc<Option<AnyBurnTensor<B, D_OUT>>> respectively
///
/// Using Box<dyn Any + Send + Sync> to avoid adding generic parameters to this struct.
/// This is (probably) safe because InferenceRequest is only sent to the actor from the coordinator layer, both of which are unavailable to the user.
pub(crate) struct InferenceRequest {
    pub(crate) observation: Box<dyn Any + Send + Sync>,
    pub(crate) mask: Box<dyn Any + Send + Sync>,
    pub(crate) reward: f32,
    pub(crate) reply_to: oneshot::Sender<Arc<RelayRLAction>>,
}

pub(crate) struct BatchedInferenceRequest {
    pub(crate) env_ids: Vec<Uuid>,
    pub(crate) env_labels: Vec<String>,
    pub(crate) observations: Box<dyn Any + Send + Sync>,
    pub(crate) masks: Box<dyn Any + Send + Sync>,
    pub(crate) rewards: Vec<f32>,
    pub(crate) reply_to: oneshot::Sender<Vec<RelayRLAction>>,
}