fix: type checking of dataclasses

[dluo96]

Is your feature request related to a problem? Please describe

At multiple places in the code, a workaround is used to check typing of chex dataclasses. This involves doing a conditional import using if TYPE_CHECKING - we would like to avoid this if possible. It is related to this issue.

Describe the solution you'd like

We would a solution which avoids doing the conditional import.

Alternatives considered

Use a built-in dataclass instead of a chex dataclass. Example implementation where mypy doesn't complain:

@dataclasses.dataclass(init=False)
class TimeStep(Generic[Observation]):
    step_type: StepType
    reward: Array
    discount: Array
    observation: Observation
    extras: Optional[Dict]

    def __init__(self, step_type: StepType, reward: Array, discount: Array, observation: Observation,
                 extras: Optional[Dict] = None):
        self.step_type = step_type
        self.reward = reward
        self.discount = discount
        self.observation = observation
        self.extras = extras

    def first(self) -> Array:
        return self.step_type == StepType.FIRST

    def mid(self) -> Array:
        return self.step_type == StepType.MID

    def last(self) -> Array:
        return self.step_type == StepType.LAST