import ray
from ray.rllib.policy.policy import Policy as BasePolicy
from ray.rllib.agents.trainer_template import build_trainer


class Policy(BasePolicy):
    def compute_actions(self, obs_batch, **kwargs):
        actions = [self.action_space.sample() for _ in obs_batch]
        return actions, [], {}

    def get_weights(self):
        pass
    
    def set_weights(self, weights):
        pass

    def learn_on_batch(self, sample_batch):
        pass


trainer = build_trainer(
    name='foo',
    default_policy=Policy,
)

ray.init()
ray.tune.run(
    trainer,
    config={
        'framework': 'jax',
        'env': 'FrozenLake-v0',
    },
    stop={'training_iteration': 1}
)