def actions(self, state, context=None):
    """Returns the next action for the state.

    Args:
      state: A [-1, num_state_dims] tensor representing a state.