forked from huggingface/diffusers
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathrun_diffuser_value_guided.py
More file actions
69 lines (56 loc) 路 1.71 KB
/
Copy pathrun_diffuser_value_guided.py
File metadata and controls
69 lines (56 loc) 路 1.71 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import d4rl # noqa
import gym
import tqdm
from diffusers import DiffusionPipeline
config = dict(
n_samples=64,
horizon=32,
num_inference_steps=20,
n_guide_steps=2,
scale_grad_by_std=True,
scale=0.1,
eta=0.0,
t_grad_cutoff=2,
device="cpu",
)
def _run():
env_name = "hopper-medium-v2"
env = gym.make(env_name)
pipeline = DiffusionPipeline.from_pretrained(
"bglick13/hopper-medium-v2-value-function-hor32",
env=env,
custom_pipeline="/Users/bglickenhaus/Documents/diffusers/examples/community",
)
# add a batch dimension and repeat for multiple samples
# [ observation_dim ] --> [ n_samples x observation_dim ]
env.seed(0)
obs = env.reset()
total_reward = 0
total_score = 0
T = 1000
rollout = [obs.copy()]
try:
for t in tqdm.tqdm(range(T)):
# 1. Call the policy
# normalize observations for forward passes
denorm_actions = pipeline(obs, planning_horizon=32)
# execute action in environment
next_observation, reward, terminal, _ = env.step(denorm_actions)
score = env.get_normalized_score(total_reward)
# update return
total_reward += reward
total_score += score
print(
f"Step: {t}, Reward: {reward}, Total Reward: {total_reward}, Score: {score}, Total Score:"
f" {total_score}"
)
# save observations for rendering
rollout.append(next_observation.copy())
obs = next_observation
except KeyboardInterrupt:
pass
print(f"Total reward: {total_reward}")
def run():
_run()
if __name__ == "__main__":
run()