From fa707c81e5dd1dc75bc47296dd0fd325ab68230b Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 6 Jan 2024 00:39:55 -0500 Subject: [PATCH] move beautiful cartpole action sampling inside jit (#3028) tested by getting 3 full scores in a row --- examples/beautiful_cartpole.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/examples/beautiful_cartpole.py b/examples/beautiful_cartpole.py index df7079aaf2..1b6306f20e 100644 --- a/examples/beautiful_cartpole.py +++ b/examples/beautiful_cartpole.py @@ -58,10 +58,10 @@ if __name__ == "__main__": return action_loss.realize(), entropy_loss.realize(), critic_loss.realize() @TinyJit - def get_action_dist(obs:Tensor) -> Tensor: + def get_action(obs:Tensor) -> Tensor: # TODO: with no_grad Tensor.no_grad = True - ret = model(obs)[0].exp().realize() + ret = model(obs)[0].exp().multinomial().realize() Tensor.no_grad = False return ret @@ -70,16 +70,15 @@ if __name__ == "__main__": st, steps = time.perf_counter(), 0 Xn, An, Rn = [], [], [] for i in (t:=trange(40)): - get_action_dist.reset() # NOTE: if you don't reset the jit here it captures the wrong model on the first run through + get_action.reset() # NOTE: if you don't reset the jit here it captures the wrong model on the first run through obs:np.ndarray = env.reset()[0] rews, terminated, truncated = [], False, False # NOTE: we don't want to early stop since then the rewards are wrong for the last episode while not terminated and not truncated: # pick actions - # TODO: move the multinomial into jitted tinygrad when JIT rand works # TODO: what's the temperature here? - act = get_action_dist(Tensor(obs)).multinomial().item() + act = get_action(Tensor(obs)).item() # save this state action pair # TODO: don't use np.copy here on the CPU, what's the tinygrad way to do this and keep on device? need __setitem__ assignment