Trying attention method

2026-04-26 22:28:43 +01:00
parent 57b1735e1a
commit a2363d882f
4 changed files with 448 additions and 12 deletions
@@ -78,9 +78,10 @@ def make_eval_env(n_sheep, seed, max_steps, reward_cfg=None):


 def run_and_record(model, vn_template, n_sheep, max_steps,
-                   reward_cfg=None, seed=42):
+                   reward_cfg=None, seed=42, make_env_fn=None):
    """Run one deterministic episode and return full trajectory history."""
-    raw = DummyVecEnv([make_eval_env(n_sheep, seed, max_steps, reward_cfg)])
+    _factory = make_env_fn or make_eval_env
+    raw = DummyVecEnv([_factory(n_sheep, seed, max_steps, reward_cfg)])
    vn = VecNormalize(raw, norm_obs=True, norm_reward=False, training=False)
    vn.obs_rms = deepcopy(vn_template.obs_rms)
    vn.ret_rms = deepcopy(vn_template.ret_rms)