File size: 3,130 Bytes
fa70ae5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
# Core game engine that manages:
# 1. Loading/parsing Wikipedia data
# 2. Game state management
# 3. Navigation between articles

import json
import random

class WikiRunEnvironment:
    def __init__(self, wiki_data_path):
        """Initialize with path to Wikipedia data"""
        self.wiki_data = self._load_wiki_data(wiki_data_path)
        self.current_article = None
        self.target_article = None
        self.path_taken = []
        self.steps = 0
        
    def _load_wiki_data(self, path):
        """Load Wikipedia data from JSON file"""
        print(f"Loading wiki data from {path}...")
        with open(path, 'r', encoding='utf-8') as f:
            wiki_data = json.load(f)
        print(f"Loaded {len(wiki_data)} articles")
        return wiki_data
    
    def reset(self, start_article=None, target_article=None):
        """Reset the environment with new start/target articles"""
        if start_article is None or target_article is None:
            # Choose random articles if not specified
            available_articles = list(self.wiki_data.keys())
            
            if start_article is None:
                start_article = random.choice(available_articles)
                
            if target_article is None:
                # Ensure target is different from start
                remaining = [a for a in available_articles if a != start_article]
                target_article = random.choice(remaining)
        
        self.current_article = start_article
        self.target_article = target_article
        self.path_taken = [start_article]
        self.steps = 0
        
        return self.get_current_state()
    
    def get_current_state(self):
        """Get the current state of the environment"""
        if self.current_article is None:
            return None
            
        current = self.wiki_data.get(self.current_article, {})
        return {
            'current_article': self.current_article,
            'target_article': self.target_article,
            'article_text': current.get('text', ''),
            'available_links': current.get('links', []),
            'steps_taken': self.steps,
            'path_taken': self.path_taken,
            'is_complete': self.current_article == self.target_article
        }
    
    def step(self, next_article):
        """Take a step to the next article"""
        if self.current_article is None:
            return None, "Game not initialized. Call reset() first."
            
        current = self.wiki_data.get(self.current_article, {})
        available_links = current.get('links', [])
        
        if next_article not in available_links:
            return self.get_current_state(), f"Invalid link: {next_article} not in available links"
            
        # Update state
        self.current_article = next_article
        self.path_taken.append(next_article)
        self.steps += 1
        
        # Check if we've reached the target
        is_complete = self.current_article == self.target_article
        
        return self.get_current_state(), "Target reached!" if is_complete else ""