action mask for DiscreteSACPolicy about tianshou HOT 4 CLOSED

# pip3 install torchrl
from torchrl.modules import MaskedCategorical

class MyDiscreteSACPolicy(DiscreteSACPolicy):
    def forward(  # type: ignore
            self,
            batch: Batch,
            state: Optional[Union[dict, Batch, np.ndarray]] = None,
            input: str = "obs",
            **kwargs: Any,
    ) -> Batch:
        obs = batch[input]
        if hasattr(batch.info, 'mask'):
            action_mask = getattr(batch.info, 'mask')
            logits, hidden = self.actor(obs, state=state, info=batch.info)
            mask = torch.tensor(action_mask)
            dist = MaskedCategorical(logits=logits, mask=mask)
        else:
            # mask = None
            logits, hidden = self.actor(obs, state=state, info=batch.info)
            dist = Categorical(logits=logits)

        if self._deterministic_eval and not self.training:
            act = dist.logits.argmax(axis=-1)
            # act = logits.argmax(axis=-1)
            # print(logits, mask, act)
        else:
            act = dist.sample()
     return Batch(act=act, state=hidden, dist=dist)
	 
	 
class MyEnv:
    def reset(self, seed=None, options=None):
        ...
        info = {'mask': action_mask}
        obs = state_arr
        return obs, info

    def step(self, action):
	...
	info = {'mask': action_mask}
        obs = state_arr
        return obs, reward, terminated, truncated, info