(state, us)
| 113 | |
| 114 | |
| 115 | def render_us(state, us): |
| 116 | rollout = [] |
| 117 | rew_sum = 0.0 |
| 118 | for i in range(Hsample): |
| 119 | rollout.append(state.pipeline_state) |
| 120 | state = step_env_jit(state, us[i]) |
| 121 | rew_sum += state.reward |
| 122 | return rollout |
| 123 | |
| 124 | |
| 125 | rng = jax.random.PRNGKey(0) |
no outgoing calls
no test coverage detected