Skip to content

Mighty exploration policy

mighty.mighty_exploration.mighty_exploration_policy #

Mighty Exploration Policy.

MightyExplorationPolicy #

MightyExplorationPolicy(algo, model, discrete=False)

Generic Exploration Policy Interface.

Now supports
  • Discrete: model(state) → logits → Categorical
  • Continuous (squashed-Gaussian): model(state) → (action, z, mean, log_std)
  • Continuous (Standard PPO): model(state) → (action, mean, log_std)
  • Continuous (legacy): model(state) → (mean, std)

:param discrete: True if action-space is discrete

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def __init__(
    self,
    algo,
    model,
    discrete=False,
) -> None:
    """
    :param algo:    Algorithm name (e.g. "ppo", "sac", etc.)
    :param model:   The policy network (any nn.Module)
    :param discrete: True if action-space is discrete
    """
    self.rng = np.random.default_rng()
    self.algo = algo
    self.model = model
    self.discrete = discrete

    # Check which action sampling to use
    if self.algo == "q":
        self.sample_action = self.sample_func_q
    else:
        self.sample_action = self.sample_func_logits

__call__ #

__call__(
    s, return_logp=False, metrics=None, evaluate=False
)

Get action.

:param s: state :param return_logp: return logprobs :param metrics: current metric dict :param eval: eval mode :return: action or (action, logprobs)

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __call__(self, s, return_logp=False, metrics=None, evaluate=False):
    """Get action.

    :param s: state
    :param return_logp: return logprobs
    :param metrics: current metric dict
    :param eval: eval mode
    :return: action or (action, logprobs)
    """
    if metrics is None:
        metrics = {}
    if evaluate:
        action, logprobs = self.sample_action(s)
        output = (action, logprobs) if return_logp else action
    else:
        output = self.explore(s, return_logp, metrics)

    return output

explore #

explore(s, return_logp, metrics=None)

Explore.

:param s: state :param return_logp: return logprobs :param _: not used :return: action or (action, logprobs)

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
171
172
173
174
175
176
177
178
179
180
def explore(self, s, return_logp, metrics=None):
    """Explore.

    :param s: state
    :param return_logp: return logprobs
    :param _: not used
    :return: action or (action, logprobs)
    """
    action, logprobs = self.explore_func(s)
    return (action, logprobs) if return_logp else action

explore_func #

explore_func(s)

Explore function.

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
182
183
184
def explore_func(self, s):
    """Explore function."""
    raise NotImplementedError

sample_func_logits #

sample_func_logits(state_array)

state_np: np.ndarray of shape [batch, obs_dim] Returns: (action_tensor, log_prob_tensor)

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def sample_func_logits(self, state_array):
    """
    state_np: np.ndarray of shape [batch, obs_dim]
    Returns: (action_tensor, log_prob_tensor)
    """
    state = torch.as_tensor(state_array, dtype=torch.float32)

    # ─── Discrete action branch ─────────────────────────────────────────
    if self.discrete:
        logits = self.model(state)  # [batch, n_actions]
        dist = Categorical(logits=logits)
        action = dist.sample()  # [batch]
        log_prob = dist.log_prob(action)  # [batch]
        return action.detach().cpu().numpy(), log_prob

    # ─── Continuous action branches ─────────────────────────────────────
    out = self.model(state)

    # NEW: Handle 3-tuple (Standard PPO)
    if isinstance(out, tuple) and len(out) == 3:
        action, mean, log_std = out
        std = torch.exp(log_std)
        dist = Normal(mean, std)
        log_prob = dist.log_prob(action).sum(dim=-1)  # Direct log prob
        return action.detach().cpu().numpy(), log_prob

    # ─── Continuous squashed‐Gaussian (4‐tuple) ──────────────────────────
    elif isinstance(out, tuple) and len(out) == 4:
        action = out[0]  # [batch, action_dim]
        log_prob = sample_nondeterministic_logprobs(
            z=out[1], mean=out[2], log_std=out[3], sac=self.ago == "sac"
        )
        return action.detach().cpu().numpy(), log_prob

    # ─── Legacy continuous branch (model returns (mean, std)) ────────────
    elif isinstance(out, tuple) and len(out) == 2:
        mean, std = out  # both [batch, action_dim]
        dist = Normal(mean, std)
        z = dist.rsample()  # [batch, action_dim]
        action = torch.tanh(z)  # [batch, action_dim]

        # 3a) log_pz = ∑ᵢ log N(zᵢ; μᵢ, σᵢ)
        log_pz = dist.log_prob(z).sum(dim=-1)  # [batch]

        # 3b) tanh‐correction
        eps = 1e-6
        log_correction = torch.log(1.0 - action.pow(2) + eps).sum(dim=-1)  # [batch]

        log_prob = log_pz - log_correction  # [batch]
        return action.detach().cpu().numpy(), log_prob

    # ─── Fallback: if model(state) returns a Distribution ────────────────
    elif isinstance(out, torch.distributions.Distribution):
        dist = out  # user returned a Distribution
        action = dist.sample()  # [batch]
        log_prob = dist.log_prob(action)  # [batch]
        return action.detach().cpu().numpy(), log_prob

    # ─── Otherwise, we don't know how to sample ─────────────────────────
    else:
        raise RuntimeError(
            "MightyExplorationPolicy: cannot interpret model(state) output of type "
            f"{type(out)}"
        )

sample_func_q #

sample_func_q(state_array)
Q-learning branch

• state_np: np.ndarray of shape [batch, obs_dim] • model(state) returns Q-values: tensor [batch, n_actions]

We choose action = argmax(Q), and also return the full Q‐vector.

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
74
75
76
77
78
79
80
81
82
83
84
85
def sample_func_q(self, state_array):
    """
    Q-learning branch:
      • state_np: np.ndarray of shape [batch, obs_dim]
      • model(state) returns Q-values: tensor [batch, n_actions]
    We choose action = argmax(Q), and also return the full Q‐vector.
    """
    state = torch.as_tensor(state_array, dtype=torch.float32)
    qs = self.model(state)  # [batch, n_actions]
    # Choose greedy action
    action = torch.argmax(qs, dim=1)  # [batch]
    return action.detach().cpu().numpy(), qs  # action_np, Q‐vector

seed #

seed(seed: int) -> None

Set the random seed for reproducibility.

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
70
71
72
def seed(self, seed: int) -> None:
    """Set the random seed for reproducibility."""
    self.rng = np.random.default_rng(seed)

sample_nondeterministic_logprobs #

sample_nondeterministic_logprobs(
    z: Tensor,
    mean: Tensor,
    log_std: Tensor,
    sac: bool = False,
) -> Tensor

Compute log-prob of a Gaussian sample z ~ N(mean, exp(log_std)), and if sac=True apply the tanh-squash correction to get log π(a).

Source code in mighty/mighty_exploration/mighty_exploration_policy.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def sample_nondeterministic_logprobs(
    z: torch.Tensor, mean: torch.Tensor, log_std: torch.Tensor, sac: bool = False
) -> torch.Tensor:
    """
    Compute log-prob of a Gaussian sample z ~ N(mean, exp(log_std)),
    and if sac=True apply the tanh-squash correction to get log π(a).
    """
    std = torch.exp(log_std)  # [batch, action_dim]
    dist = Normal(mean, std)
    # base Gaussian log‐prob of z
    log_pz = dist.log_prob(z).sum(dim=-1, keepdim=True)  # [batch, 1]

    if sac:
        # subtract the ∑_i log(d tanh/dz_i) = ∑ log(1 - tanh(z)^2)
        eps = 1e-4
        log_correction = torch.log(1.0 - torch.tanh(z).pow(2) + eps).sum(
            dim=-1, keepdim=True
        )  # [batch, 1]
        return log_pz - log_correction
    else:
        # PPO-style or other: no squash correction
        return log_pz