Skip to main content

JointPolicy

thrust_rl::multi_agent::joint

Trait JointPolicy 

Source 
pub trait JointPolicy<B: AutodiffBackend>: AutodiffModule<B> + Clone {
    // Required methods
    fn get_action_host_seeded(
        &self,
        obs: Tensor<B, 2>,
        rng: &mut StdRng,
    ) -> (Vec<i64>, Vec<f32>, Vec<f32>);
    fn evaluate_actions_joint(
        &self,
        obs: Tensor<B, 2>,
        actions: Tensor<B, 2, Int>,
    ) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>);
    fn encoder_features_joint(&self, obs: Tensor<B, 2>) -> Tensor<B, 2>;
    fn action_dims_joint(&self) -> Vec<i64>;

    // Provided method
    fn get_actions_host_seeded_batched(
        &self,
        obs: Tensor<B, 2>,
        rng: &mut StdRng,
    ) -> (Vec<i64>, Vec<f32>, Vec<f32>) { ... }
}
Expand description

Capabilities a policy must expose to participate in JointMultiAgentTrainer.

The trait pins exactly the surface the trainer needs:

  • get_action_host_seeded — rollout-time sampling. Returns (actions_per_dim, log_probs, values) on the host so the trainer can build the rollout buffer without tying it to a particular backend tensor. Takes the trainer-owned StdRng so PsroConfig::seed / NfspConfig::seed produce bit-identical rollouts (issue #114).
  • evaluate_actions — re-evaluate the current policy on previously sampled actions to compute updated log-probs / entropy / value for the PPO loss; this is the only place autograd-bearing tensors are produced.
  • encoder_features — shared-trunk activations for the auxiliary loss.
  • action_dims — per-dim action cardinalities, used to size action buffers without invoking the policy.

Required Methods§

Source

fn get_action_host_seeded( &self, obs: Tensor<B, 2>, rng: &mut StdRng, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Sample actions for a single rollout step using a caller-supplied seeded RNG.

obs carries one row per environment in the rollout batch. Returns host-side (actions, log_probs, values) where:

  • actions is laid out flat per-row: actions[row * num_dims + d] is the action sampled for dim d of row row. Length is obs.dims()\[0\] * num_dims.
  • log_probs[row] is the joint log-probability summed across dims.
  • values[row] is the value estimate.

Bit-exactness: two calls with the same obs, same policy state, and same-seeded rng produce element-wise identical outputs — the load-bearing guarantee that PsroConfig::seed / NfspConfig::seed rely on after issue #114 completed plumbing the trainer-owned StdRng through the rollout-time action sampler.

Source

fn evaluate_actions_joint( &self, obs: Tensor<B, 2>, actions: Tensor<B, 2, Int>, ) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>)

Re-evaluate the policy on previously-sampled actions.

actions is shape [batch, num_dims]. For scalar discrete policies (num_dims == 1) pass actions reshaped to [batch, 1]. Returns (log_probs, entropy, values) where every tensor has shape [batch].

Source

fn encoder_features_joint(&self, obs: Tensor<B, 2>) -> Tensor<B, 2>

Shared-trunk feature representation; gradients flow back into the encoder. The natural quantity to feed into cross-agent regularizers.

Shape: [batch, hidden_dim].

Source

fn action_dims_joint(&self) -> Vec<i64>

Per-dimension action cardinalities.

Provided Methods§

Source

fn get_actions_host_seeded_batched( &self, obs: Tensor<B, 2>, rng: &mut StdRng, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Batched seeded sampler: obs carries N rows scored through this one policy in a single forward, returning N actions.

§When this helps (and when it does not)

This eliminates per-call batch-1 forward overhead on the NdArray backend wherever many observations are scored through the same model in one step — e.g. the constant-obs marginal probe (crate::multi_agent::nfsp::NfspTrainer::action_marginal_for), or a future parallel-env rollout where one frozen policy scores many env rows.

It does not apply across the per-agent forwards of NFSP’s (or PSRO’s) rollout: those num_agents forwards each run a distinct per-agent policy module, and a single batched forward applies one weight set to every row — so stacking the agents’ observations into one [N, obs_dim] tensor would be semantically wrong. NFSP runs a single shared joint episode (env.step_joint), so there is no within-model batch dimension to collapse in the rollout itself; the per-step forward count there is irreducible. See issue #235.

§Determinism

The default implementation simply calls Self::get_action_host_seeded on the full [N, obs_dim] tensor, which already draws RNG in row-major order. Implementors that override this (the concrete MLP policies do) must keep the per-row RNG draw order identical so the sampled action stream is bit-for-bit reproducible under PsroConfig::seed / NfspConfig::seed (issue #114).

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§