Add deterministic Deep CFR traversal mode
This commit is contained in:
@@ -5,6 +5,7 @@ run:
|
|||||||
max_minutes: null
|
max_minutes: null
|
||||||
device: cuda
|
device: cuda
|
||||||
use_amp: false
|
use_amp: false
|
||||||
|
deterministic: false
|
||||||
|
|
||||||
rules:
|
rules:
|
||||||
n_colors: 5
|
n_colors: 5
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ run:
|
|||||||
max_minutes: null
|
max_minutes: null
|
||||||
device: cuda
|
device: cuda
|
||||||
use_amp: false
|
use_amp: false
|
||||||
|
deterministic: false
|
||||||
|
|
||||||
rules:
|
rules:
|
||||||
n_colors: 5
|
n_colors: 5
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ class RunConfig(StrictModel):
|
|||||||
seed: int = 1
|
seed: int = 1
|
||||||
device: str = "auto"
|
device: str = "auto"
|
||||||
use_amp: bool = False
|
use_amp: bool = False
|
||||||
|
deterministic: bool = False
|
||||||
|
|
||||||
@field_validator("device")
|
@field_validator("device")
|
||||||
@classmethod
|
@classmethod
|
||||||
|
|||||||
@@ -182,12 +182,14 @@ class StrategyNetPolicy(LostCitiesPolicy):
|
|||||||
sample: bool = False,
|
sample: bool = False,
|
||||||
seed: int | None = None,
|
seed: int | None = None,
|
||||||
encoding: EncodingConfig | None = None,
|
encoding: EncodingConfig | None = None,
|
||||||
|
deterministic: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.strategy_network = strategy_network
|
self.strategy_network = strategy_network
|
||||||
self.device = torch.device(device)
|
self.device = torch.device(device)
|
||||||
self.sample = sample
|
self.sample = sample
|
||||||
self.rng = np.random.default_rng(seed)
|
self.rng = np.random.default_rng(seed)
|
||||||
self.encoding = encoding
|
self.encoding = encoding
|
||||||
|
self.deterministic = deterministic
|
||||||
self.runtime = EvalRuntimeCounters()
|
self.runtime = EvalRuntimeCounters()
|
||||||
|
|
||||||
def select_actions_batch(self, states: list[GameState]) -> list[tuple[int, float]]:
|
def select_actions_batch(self, states: list[GameState]) -> list[tuple[int, float]]:
|
||||||
@@ -209,8 +211,19 @@ class StrategyNetPolicy(LostCitiesPolicy):
|
|||||||
|
|
||||||
network_started = time.perf_counter()
|
network_started = time.perf_counter()
|
||||||
with torch.inference_mode():
|
with torch.inference_mode():
|
||||||
x = torch.as_tensor(np.stack(infos), dtype=torch.float32, device=self.device)
|
if self.deterministic:
|
||||||
logits = self.strategy_network(x)
|
logits = torch.cat(
|
||||||
|
[
|
||||||
|
self.strategy_network(
|
||||||
|
torch.as_tensor(info[None, :], dtype=torch.float32, device=self.device)
|
||||||
|
)
|
||||||
|
for info in infos
|
||||||
|
],
|
||||||
|
dim=0,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
x = torch.as_tensor(np.stack(infos), dtype=torch.float32, device=self.device)
|
||||||
|
logits = self.strategy_network(x)
|
||||||
self.runtime.policy_network_seconds += time.perf_counter() - network_started
|
self.runtime.policy_network_seconds += time.perf_counter() - network_started
|
||||||
|
|
||||||
postprocess_started = time.perf_counter()
|
postprocess_started = time.perf_counter()
|
||||||
@@ -297,6 +310,7 @@ def evaluate_strategy_network(
|
|||||||
encoding: EncodingConfig | None = None,
|
encoding: EncodingConfig | None = None,
|
||||||
batch_size: int = 64,
|
batch_size: int = 64,
|
||||||
save_games_path: str | None = None,
|
save_games_path: str | None = None,
|
||||||
|
deterministic: bool = False,
|
||||||
) -> dict[str, float | int]:
|
) -> dict[str, float | int]:
|
||||||
strategy_network.eval()
|
strategy_network.eval()
|
||||||
return _evaluate_strategy_network_with_diagnostics(
|
return _evaluate_strategy_network_with_diagnostics(
|
||||||
@@ -310,6 +324,7 @@ def evaluate_strategy_network(
|
|||||||
encoding=encoding,
|
encoding=encoding,
|
||||||
batch_size=batch_size,
|
batch_size=batch_size,
|
||||||
save_games_path=save_games_path,
|
save_games_path=save_games_path,
|
||||||
|
deterministic=deterministic,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -325,6 +340,7 @@ def _evaluate_strategy_network_with_diagnostics(
|
|||||||
encoding: EncodingConfig | None,
|
encoding: EncodingConfig | None,
|
||||||
batch_size: int,
|
batch_size: int,
|
||||||
save_games_path: str | None = None,
|
save_games_path: str | None = None,
|
||||||
|
deterministic: bool = False,
|
||||||
) -> dict[str, float | int]:
|
) -> dict[str, float | int]:
|
||||||
if games <= 0:
|
if games <= 0:
|
||||||
raise ValueError(f"games must be positive, got {games}")
|
raise ValueError(f"games must be positive, got {games}")
|
||||||
@@ -335,6 +351,7 @@ def _evaluate_strategy_network_with_diagnostics(
|
|||||||
device=device,
|
device=device,
|
||||||
seed=seed * 2,
|
seed=seed * 2,
|
||||||
encoding=encoding,
|
encoding=encoding,
|
||||||
|
deterministic=deterministic,
|
||||||
)
|
)
|
||||||
started = time.perf_counter()
|
started = time.perf_counter()
|
||||||
active_games: list[_EvalGame] = []
|
active_games: list[_EvalGame] = []
|
||||||
|
|||||||
@@ -114,6 +114,7 @@ class InterleavedTraversalConfig:
|
|||||||
opponent_policy: str
|
opponent_policy: str
|
||||||
endpoint_depth_bucket_width: int
|
endpoint_depth_bucket_width: int
|
||||||
endpoint_depth_bucket_max: int
|
endpoint_depth_bucket_max: int
|
||||||
|
deterministic: bool = False
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -182,11 +183,13 @@ class BatchedPolicy:
|
|||||||
device: torch.device,
|
device: torch.device,
|
||||||
epsilon: float,
|
epsilon: float,
|
||||||
strategy_network: torch.nn.Module | None = None,
|
strategy_network: torch.nn.Module | None = None,
|
||||||
|
deterministic: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.networks = networks
|
self.networks = networks
|
||||||
self.strategy_network = strategy_network
|
self.strategy_network = strategy_network
|
||||||
self.device = device
|
self.device = device
|
||||||
self.epsilon = epsilon
|
self.epsilon = epsilon
|
||||||
|
self.deterministic = deterministic
|
||||||
self.batch_sizes: list[int] = []
|
self.batch_sizes: list[int] = []
|
||||||
self.forward_seconds = 0.0
|
self.forward_seconds = 0.0
|
||||||
|
|
||||||
@@ -194,13 +197,27 @@ class BatchedPolicy:
|
|||||||
if not requests:
|
if not requests:
|
||||||
return []
|
return []
|
||||||
out: list[PolicyResult | None] = [None] * len(requests)
|
out: list[PolicyResult | None] = [None] * len(requests)
|
||||||
group_keys = sorted({(request.network_kind, request.player) for request in requests})
|
if self.deterministic:
|
||||||
for network_kind, player in group_keys:
|
groups = [
|
||||||
indices = [
|
(request.network_kind, request.player, [index])
|
||||||
idx
|
for index, request in enumerate(requests)
|
||||||
for idx, request in enumerate(requests)
|
|
||||||
if (request.network_kind, request.player) == (network_kind, player)
|
|
||||||
]
|
]
|
||||||
|
else:
|
||||||
|
group_keys = sorted({(request.network_kind, request.player) for request in requests})
|
||||||
|
groups = [
|
||||||
|
(
|
||||||
|
network_kind,
|
||||||
|
player,
|
||||||
|
[
|
||||||
|
idx
|
||||||
|
for idx, request in enumerate(requests)
|
||||||
|
if (request.network_kind, request.player) == (network_kind, player)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
for network_kind, player in group_keys
|
||||||
|
]
|
||||||
|
for network_kind, player, indices in groups:
|
||||||
|
indices = [idx for idx in indices]
|
||||||
states = np.stack([requests[idx].info_state for idx in indices]).astype(np.float32)
|
states = np.stack([requests[idx].info_state for idx in indices]).astype(np.float32)
|
||||||
x = torch.as_tensor(states, dtype=torch.float32, device=self.device)
|
x = torch.as_tensor(states, dtype=torch.float32, device=self.device)
|
||||||
if self.device.type == "cuda":
|
if self.device.type == "cuda":
|
||||||
@@ -605,6 +622,8 @@ def run_interleaved_traversal_batch(
|
|||||||
seed: int,
|
seed: int,
|
||||||
interleave_width: int,
|
interleave_width: int,
|
||||||
interleave_max_batch: int,
|
interleave_max_batch: int,
|
||||||
|
traversal_start_index: int = 0,
|
||||||
|
deterministic: bool = False,
|
||||||
) -> tuple[TraversalStats, list[TrainingSample], list[TrainingSample], dict[str, float | int]]:
|
) -> tuple[TraversalStats, list[TrainingSample], list[TrainingSample], dict[str, float | int]]:
|
||||||
cfg = InterleavedTraversalConfig(
|
cfg = InterleavedTraversalConfig(
|
||||||
action_size=action_size,
|
action_size=action_size,
|
||||||
@@ -620,14 +639,18 @@ def run_interleaved_traversal_batch(
|
|||||||
opponent_policy=opponent_policy,
|
opponent_policy=opponent_policy,
|
||||||
endpoint_depth_bucket_width=endpoint_depth_bucket_width,
|
endpoint_depth_bucket_width=endpoint_depth_bucket_width,
|
||||||
endpoint_depth_bucket_max=endpoint_depth_bucket_max,
|
endpoint_depth_bucket_max=endpoint_depth_bucket_max,
|
||||||
|
deterministic=deterministic,
|
||||||
)
|
)
|
||||||
states = [GameState.new_game(game_config, seed=game_seed) for game_seed in seeds]
|
states = [GameState.new_game(game_config, seed=game_seed) for game_seed in seeds]
|
||||||
rng_seeds = [int(seed) + index * 1_000_003 for index in range(len(seeds))]
|
rng_seeds = [
|
||||||
|
int(seed) + (int(traversal_start_index) + index) * 1_000_003 for index in range(len(seeds))
|
||||||
|
]
|
||||||
policy = BatchedPolicy(
|
policy = BatchedPolicy(
|
||||||
advantage_networks,
|
advantage_networks,
|
||||||
device=device,
|
device=device,
|
||||||
epsilon=cfg.epsilon,
|
epsilon=cfg.epsilon,
|
||||||
strategy_network=strategy_network,
|
strategy_network=strategy_network,
|
||||||
|
deterministic=cfg.deterministic,
|
||||||
)
|
)
|
||||||
scheduler = InterleavedTraversalScheduler(cfg, policy)
|
scheduler = InterleavedTraversalScheduler(cfg, policy)
|
||||||
_values, _rng_out, stats_rows, sample_rows, batch_sizes = scheduler.run(
|
_values, _rng_out, stats_rows, sample_rows, batch_sizes = scheduler.run(
|
||||||
|
|||||||
@@ -69,9 +69,22 @@ class EvaluationWorkerJob:
|
|||||||
device: str
|
device: str
|
||||||
max_steps: int
|
max_steps: int
|
||||||
batch_size: int
|
batch_size: int
|
||||||
|
deterministic: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
def configure_torch_determinism(enabled: bool) -> None:
|
||||||
|
torch.use_deterministic_algorithms(bool(enabled))
|
||||||
|
if not enabled:
|
||||||
|
return
|
||||||
|
if torch.backends.cudnn.is_available():
|
||||||
|
torch.backends.cudnn.benchmark = False
|
||||||
|
torch.backends.cudnn.allow_tf32 = False
|
||||||
|
if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
|
||||||
|
torch.backends.cuda.matmul.allow_tf32 = False
|
||||||
|
|
||||||
|
|
||||||
def run_evaluation_worker(job: EvaluationWorkerJob) -> tuple[str, dict[str, float | int]]:
|
def run_evaluation_worker(job: EvaluationWorkerJob) -> tuple[str, dict[str, float | int]]:
|
||||||
|
configure_torch_determinism(job.deterministic)
|
||||||
device = _resolve_torch_device(job.device)
|
device = _resolve_torch_device(job.device)
|
||||||
game_config = LostCitiesConfig(**job.game_config)
|
game_config = LostCitiesConfig(**job.game_config)
|
||||||
network_config = NetworkConfig.model_validate(job.network_config)
|
network_config = NetworkConfig.model_validate(job.network_config)
|
||||||
@@ -89,6 +102,7 @@ def run_evaluation_worker(job: EvaluationWorkerJob) -> tuple[str, dict[str, floa
|
|||||||
max_steps=job.max_steps,
|
max_steps=job.max_steps,
|
||||||
encoding=encoding_config,
|
encoding=encoding_config,
|
||||||
batch_size=job.batch_size,
|
batch_size=job.batch_size,
|
||||||
|
deterministic=job.deterministic,
|
||||||
)
|
)
|
||||||
return job.opponent, result
|
return job.opponent, result
|
||||||
|
|
||||||
@@ -201,6 +215,7 @@ class DeepCFRTrainer:
|
|||||||
extra_trackers: list[RunTracker] | None = None,
|
extra_trackers: list[RunTracker] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.config = config or DeepCFRConfig()
|
self.config = config or DeepCFRConfig()
|
||||||
|
configure_torch_determinism(self.config.run.deterministic)
|
||||||
self.game_config = game_config or self.config.rules.to_lost_cities_config(
|
self.game_config = game_config or self.config.rules.to_lost_cities_config(
|
||||||
seed=self.config.run.seed
|
seed=self.config.run.seed
|
||||||
)
|
)
|
||||||
@@ -327,9 +342,11 @@ class DeepCFRTrainer:
|
|||||||
self._ensure_inference_server()
|
self._ensure_inference_server()
|
||||||
self._maybe_sync_inference_server(iteration)
|
self._maybe_sync_inference_server(iteration)
|
||||||
traversal_started = time.perf_counter()
|
traversal_started = time.perf_counter()
|
||||||
|
worker_count = self.config.traversal.resolved_num_workers()
|
||||||
if (
|
if (
|
||||||
self.config.traversal.resolved_num_workers() > 1
|
worker_count > 1
|
||||||
or self.config.traversal.inference_backend == "server"
|
or self.config.traversal.inference_backend == "server"
|
||||||
|
or (self.config.run.deterministic and worker_count > 0)
|
||||||
):
|
):
|
||||||
total_stats = self._run_traversals_parallel(iteration)
|
total_stats = self._run_traversals_parallel(iteration)
|
||||||
else:
|
else:
|
||||||
@@ -435,6 +452,7 @@ class DeepCFRTrainer:
|
|||||||
seed=self.config.run.seed + iteration * 1_000_003 + player,
|
seed=self.config.run.seed + iteration * 1_000_003 + player,
|
||||||
interleave_width=self.config.traversal.interleave_width,
|
interleave_width=self.config.traversal.interleave_width,
|
||||||
interleave_max_batch=self.config.traversal.interleave_max_batch,
|
interleave_max_batch=self.config.traversal.interleave_max_batch,
|
||||||
|
deterministic=self.config.run.deterministic,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
self._record_interleaved_metrics(runtime_metrics)
|
self._record_interleaved_metrics(runtime_metrics)
|
||||||
@@ -545,6 +563,7 @@ class DeepCFRTrainer:
|
|||||||
for _, batch in zip(range(in_flight_limit), batch_iter, strict=False)
|
for _, batch in zip(range(in_flight_limit), batch_iter, strict=False)
|
||||||
}
|
}
|
||||||
completed_batches = 0
|
completed_batches = 0
|
||||||
|
completed_results = []
|
||||||
while futures:
|
while futures:
|
||||||
done, futures = wait(futures, timeout=5.0, return_when=FIRST_COMPLETED)
|
done, futures = wait(futures, timeout=5.0, return_when=FIRST_COMPLETED)
|
||||||
if not done:
|
if not done:
|
||||||
@@ -557,17 +576,8 @@ class DeepCFRTrainer:
|
|||||||
continue
|
continue
|
||||||
for future in done:
|
for future in done:
|
||||||
result = future.result()
|
result = future.result()
|
||||||
|
completed_results.append(result)
|
||||||
completed_batches += 1
|
completed_batches += 1
|
||||||
total_stats.accumulate(result.stats)
|
|
||||||
self._record_interleaved_metrics(result.runtime_metrics)
|
|
||||||
memory_add_started = time.perf_counter()
|
|
||||||
self._add_advantage_samples(result.advantage_samples)
|
|
||||||
self.strategy_memory.add_many(result.strategy_samples, self.rng)
|
|
||||||
self._runtime_metrics["time/memory_add_seconds"] = (
|
|
||||||
float(self._runtime_metrics.get("time/memory_add_seconds", 0.0))
|
|
||||||
+ time.perf_counter()
|
|
||||||
- memory_add_started
|
|
||||||
)
|
|
||||||
progress_nodes += result.stats.nodes
|
progress_nodes += result.stats.nodes
|
||||||
progress_traversals += result.traversals
|
progress_traversals += result.traversals
|
||||||
if next_progress_at is not None and progress_traversals >= next_progress_at:
|
if next_progress_at is not None and progress_traversals >= next_progress_at:
|
||||||
@@ -587,6 +597,20 @@ class DeepCFRTrainer:
|
|||||||
next_batch = next(batch_iter, None)
|
next_batch = next(batch_iter, None)
|
||||||
if next_batch is not None:
|
if next_batch is not None:
|
||||||
futures.add(executor.submit(run_traversal_worker_batch, next_batch))
|
futures.add(executor.submit(run_traversal_worker_batch, next_batch))
|
||||||
|
memory_add_started = time.perf_counter()
|
||||||
|
for result in sorted(
|
||||||
|
completed_results,
|
||||||
|
key=lambda item: (item.player, item.traversal_start_index, item.batch_index),
|
||||||
|
):
|
||||||
|
total_stats.accumulate(result.stats)
|
||||||
|
self._record_interleaved_metrics(result.runtime_metrics)
|
||||||
|
self._add_advantage_samples(result.advantage_samples)
|
||||||
|
self.strategy_memory.add_many(result.strategy_samples, self.rng)
|
||||||
|
self._runtime_metrics["time/memory_add_seconds"] = (
|
||||||
|
float(self._runtime_metrics.get("time/memory_add_seconds", 0.0))
|
||||||
|
+ time.perf_counter()
|
||||||
|
- memory_add_started
|
||||||
|
)
|
||||||
if (
|
if (
|
||||||
self.config.traversal.inference_backend == "server"
|
self.config.traversal.inference_backend == "server"
|
||||||
and self._inference_server is not None
|
and self._inference_server is not None
|
||||||
@@ -647,8 +671,10 @@ class DeepCFRTrainer:
|
|||||||
chunk = seeds[start : start + chunk_size]
|
chunk = seeds[start : start + chunk_size]
|
||||||
batches.append(
|
batches.append(
|
||||||
TraversalWorkerBatch(
|
TraversalWorkerBatch(
|
||||||
|
batch_index=batch_index,
|
||||||
player=player,
|
player=player,
|
||||||
iteration=iteration,
|
iteration=iteration,
|
||||||
|
traversal_start_index=start,
|
||||||
seeds=chunk,
|
seeds=chunk,
|
||||||
config=self.config.to_dict(),
|
config=self.config.to_dict(),
|
||||||
game_config=self.game_config.to_snapshot(),
|
game_config=self.game_config.to_snapshot(),
|
||||||
@@ -656,7 +682,7 @@ class DeepCFRTrainer:
|
|||||||
action_size=self.action_size,
|
action_size=self.action_size,
|
||||||
advantage_networks=network_payloads,
|
advantage_networks=network_payloads,
|
||||||
league_advantage_networks=league_payloads,
|
league_advantage_networks=league_payloads,
|
||||||
worker_seed=self.config.run.seed + iteration * 1_000_003 + batch_index,
|
worker_seed=self.config.run.seed + iteration * 1_000_003 + player,
|
||||||
strategy_network=strategy_payload,
|
strategy_network=strategy_payload,
|
||||||
inference_handles=None,
|
inference_handles=None,
|
||||||
)
|
)
|
||||||
@@ -884,6 +910,7 @@ class DeepCFRTrainer:
|
|||||||
max_steps=self.config.evaluation.max_steps,
|
max_steps=self.config.evaluation.max_steps,
|
||||||
encoding=self.config.encoding,
|
encoding=self.config.encoding,
|
||||||
batch_size=self.config.evaluation.batch_size,
|
batch_size=self.config.evaluation.batch_size,
|
||||||
|
deterministic=self.config.run.deterministic,
|
||||||
)
|
)
|
||||||
for key, value in result.items():
|
for key, value in result.items():
|
||||||
results[f"eval/{opponent}/{key}"] = value
|
results[f"eval/{opponent}/{key}"] = value
|
||||||
|
|||||||
@@ -43,6 +43,17 @@ def _configure_worker_torch_threads() -> None:
|
|||||||
_TORCH_THREADS_CONFIGURED = True
|
_TORCH_THREADS_CONFIGURED = True
|
||||||
|
|
||||||
|
|
||||||
|
def _configure_worker_torch_determinism(enabled: bool) -> None:
|
||||||
|
torch.use_deterministic_algorithms(bool(enabled))
|
||||||
|
if not enabled:
|
||||||
|
return
|
||||||
|
if torch.backends.cudnn.is_available():
|
||||||
|
torch.backends.cudnn.benchmark = False
|
||||||
|
torch.backends.cudnn.allow_tf32 = False
|
||||||
|
if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
|
||||||
|
torch.backends.cuda.matmul.allow_tf32 = False
|
||||||
|
|
||||||
|
|
||||||
def initialize_traversal_worker(inference_handles: InferenceClientHandles | None = None) -> None:
|
def initialize_traversal_worker(inference_handles: InferenceClientHandles | None = None) -> None:
|
||||||
global _INFERENCE_HANDLES
|
global _INFERENCE_HANDLES
|
||||||
_INFERENCE_HANDLES = inference_handles
|
_INFERENCE_HANDLES = inference_handles
|
||||||
@@ -51,8 +62,10 @@ def initialize_traversal_worker(inference_handles: InferenceClientHandles | None
|
|||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class TraversalWorkerBatch:
|
class TraversalWorkerBatch:
|
||||||
|
batch_index: int
|
||||||
player: int
|
player: int
|
||||||
iteration: int
|
iteration: int
|
||||||
|
traversal_start_index: int
|
||||||
seeds: list[int]
|
seeds: list[int]
|
||||||
config: dict[str, Any]
|
config: dict[str, Any]
|
||||||
game_config: dict[str, Any]
|
game_config: dict[str, Any]
|
||||||
@@ -67,7 +80,9 @@ class TraversalWorkerBatch:
|
|||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class TraversalWorkerResult:
|
class TraversalWorkerResult:
|
||||||
|
batch_index: int
|
||||||
player: int
|
player: int
|
||||||
|
traversal_start_index: int
|
||||||
stats: TraversalStats
|
stats: TraversalStats
|
||||||
advantage_samples: list[TrainingSample]
|
advantage_samples: list[TrainingSample]
|
||||||
strategy_samples: list[TrainingSample]
|
strategy_samples: list[TrainingSample]
|
||||||
@@ -79,6 +94,7 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
|
|||||||
_configure_worker_torch_threads()
|
_configure_worker_torch_threads()
|
||||||
|
|
||||||
cfg = config_from_dict(batch.config)
|
cfg = config_from_dict(batch.config)
|
||||||
|
_configure_worker_torch_determinism(cfg.run.deterministic)
|
||||||
device = torch.device("cpu")
|
device = torch.device("cpu")
|
||||||
client: InferenceClient | None = None
|
client: InferenceClient | None = None
|
||||||
try:
|
try:
|
||||||
@@ -176,6 +192,8 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
|
|||||||
seed=batch.worker_seed,
|
seed=batch.worker_seed,
|
||||||
interleave_width=cfg.traversal.interleave_width,
|
interleave_width=cfg.traversal.interleave_width,
|
||||||
interleave_max_batch=cfg.traversal.interleave_max_batch,
|
interleave_max_batch=cfg.traversal.interleave_max_batch,
|
||||||
|
traversal_start_index=batch.traversal_start_index,
|
||||||
|
deterministic=cfg.run.deterministic,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
@@ -221,7 +239,9 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
|
|||||||
if client is not None:
|
if client is not None:
|
||||||
client.close()
|
client.close()
|
||||||
return TraversalWorkerResult(
|
return TraversalWorkerResult(
|
||||||
|
batch_index=batch.batch_index,
|
||||||
player=batch.player,
|
player=batch.player,
|
||||||
|
traversal_start_index=batch.traversal_start_index,
|
||||||
stats=total_stats,
|
stats=total_stats,
|
||||||
advantage_samples=advantage_samples,
|
advantage_samples=advantage_samples,
|
||||||
strategy_samples=strategy_samples,
|
strategy_samples=strategy_samples,
|
||||||
|
|||||||
Reference in New Issue
Block a user