SmartPlay is a benchmark for Large Language Models (LLMs). Uses a variety of games to test various important LLM capabilities as agents. SmartPlay is designed to be easy to use, and to support future development of LLMs.
import os
os.environ["MINEDOJO_HEADLESS"]="1"
import argparse
import numpy as np
from tqdm import tqdm
import gym
import smartplay

parser = argparse.ArgumentParser()
parser.add_argument('--llm_name', type=str, default='gpt-4', help='Name of the LLM')
parser.add_argument('--env_names', type=str, default=None, help='Comma separated list of environments to run')

args = parser.parse_args()

if args.env_names is None:
    args.env_names = ','.join(smartplay.benchmark_games_v0)

LLM_name = args.llm_name

from llm_api import get_query
query_model = get_query(LLM_name)

def compose_ingame_prompt(info, question, past_qa=[]):
    messages = [
        {"role": "system", "content" : "You’re a player trying to play the game."}
    ]
    
    if len(info['manual'])>0:
        messages.append({"role": "system", "content": info['manual']})

    if len(info['history'])>0:
        messages.append({"role": "system", "content": info['history']})

    messages.append({"role": "system", "content": "current step observation: {}".format(info['obs'])})

    if len(past_qa)>0:
        for q,a in past_qa:
            messages.append({"role": "user", "content": q})
            messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": question})

    return messages, 2 # This is the index of the history, we will truncate the history if it is too long for LLM

questions=[
        "What is the best action to take? Let's think step by step, ",
        "Choose the best executable action from the list of all actions. Write the exact chosen action."
    ]

def run(env_name):
    normalized_scores = []
    env = gym.make("smartplay:{}-v0".format(env_name))
    env_steps = env.default_steps
    num_iter = env.default_iter

    def match_act(output):
        inds = [(i, output.lower().index(act.lower())) for i, act in enumerate(env.action_list) if act.lower() in output.lower()]
        if len(inds)>0:
            # return the action with smallest index
            return sorted(inds, key=lambda x:x[1])[0][0]
        else:
            # print("LLM failed with output \"{}\", taking action 0...".format(output))
            return 0

    rewards = []
    progresses = []
    for eps in tqdm(range(num_iter), desc="Evaluating LLM {} on {}".format(LLM_name, env_name)):
        import wandb
        wandb.init(project="SmartPlay", config={"LLM": LLM_name, "env": env_name, "eps": eps, "num_iter": num_iter, "env_steps": env_steps})
        step = 0
        trajectories = []
        qa_history = []
        progress = [0]
        reward = 0
        rewards = []
        done=False

        columns=["Context", "Step", "OBS", "History", "Score", "Reward", "Total Reward"] + questions + ["Action"]
        wandb_table = wandb.Table(columns=columns)

        _, info = env.reset()
        
        while step < env_steps:

            new_row = [info['manual'], step, info['obs'], info['history'], info['score'], reward, sum(rewards)]
            wandb.log({"metric/total_reward".format(eps): sum(rewards), 
                       "metric/score".format(eps): info['score'],
                       "metric/reward".format(eps): reward,
                       })
            
            if done:
                break
            
            qa_history = []
            for question in questions:
                prompt = compose_ingame_prompt(info, question, qa_history)
                answer = query_model(*prompt)
                qa_history.append((question, answer))
                new_row.append(answer)
                answer_act = answer

            a = match_act(answer_act)
            new_row.append(env.action_list[a])
            _, reward, done, info = env.step(a)
            rewards.append(reward)
            score=info['score']

            step += 1
            wandb_table.add_data(*new_row)

        if not done:
            completion=0
        else:
            completion=info['completed']
        progresses.append(np.max(progress))
        wandb.log({"rollout/rollout".format(eps): wandb_table, 
                "final/total_reward":sum(rewards),
                "final/score":score,
                "final/normalized_score":smartplay.normalize_score(env_name, score),
                "final/completion":completion,
                "final/episodic_step":step,
                "final/eps":eps,
                })
        normalized_scores.append(smartplay.normalize_score(env_name, score))
        del wandb_table
        wandb.finish()
    return np.average(normalized_scores)

score_dict = {}
for env_name in args.env_names.split(','):
    score_dict[env_name] = run(env_name)

print("Normalized scores on each task:", score_dict)
print("Capability scores of the LLM:", smartplay.analyze_capabilities(score_dict))

examples/experiment.py
microsoft/SmartPlay
import setuptools
import pathlib


setuptools.setup(
    name='smartplay',
    version='0.0.1',
    description='A benchmarking tool for LLMs with games',
    url='',
    long_description=pathlib.Path('README.md').read_text(),
    long_description_content_type='text/markdown',
    packages=setuptools.find_packages('src'),
    package_dir={'': 'src'},
    # package_data={'crafter': ['data.yaml', 'assets/*']},
    entry_points={'console_scripts': ['smartplay=smartplay.run_gui:main']},
    install_requires=[
        'numpy', 
        'pandas',
        'pygame', 
        'connected-components-3d', 
        'gym', 
        'minedojo', 
        'imageio', 
        'pillow', 
        'opensimplex', 
        'ruamel.yaml',
        'importlib-metadata==6.6.0',
        'importlib-resources==5.12.0',
        'vgdl @ git+https://github.com/ahjwang/py-vgdl',
    ],
    include_package_data=True,
    classifiers=[
        'Intended Audience :: Science/Research',
        'License :: OSI Approved :: MIT License',
        'Programming Language :: Python :: 3',
        'Topic :: Games/Entertainment',
        'Topic :: Scientific/Engineering :: Artificial Intelligence',
    ],
)

setup.py
microsoft/SmartPlay
from gym.envs.registration import register

register(
    id = "msgr-train-v1",
    entry_point="messenger.envs:TwoEnvWrapper",
    kwargs = dict(
        stage=1,
        split_1="train_mc",
        split_2="train_sc",
        prob_env_1=0.75
    )
)

register(
    id = "msgr-train-sc-v1",
    entry_point="messenger.envs:StageOne",
    kwargs = dict(
        split="train_sc",
    )
)

register(
    id = "msgr-train-mc-v1",
    entry_point="messenger.envs:StageOne",
    kwargs = dict(
        split="train_mc",
    )
)

register(
    id = "msgr-val-v1",
    entry_point="messenger.envs:StageOne",
    kwargs = dict(
        split="val",
    )
)

register(
    id = "msgr-test-v1",
    entry_point="messenger.envs:StageOne",
    kwargs = dict(
        split="test",
    )
)

register(
    id = "msgr-train-v2",
    entry_point="messenger.envs:TwoEnvWrapper",
    kwargs = dict(
        stage=2,
        split_1="train_mc",
        split_2="train_sc",
        prob_env_1=0.75
    )
)

register(
    id = "msgr-train-sc-v2",
    entry_point="messenger.envs:StageTwo",
    kwargs = dict(
        split="train_sc",
    )
)

register(
    id = "msgr-train-mc-v2",
    entry_point="messenger.envs:StageTwo",
    kwargs = dict(
        split="train_mc",
    )
)

register(
    id = "msgr-val-v2",
    entry_point="messenger.envs:StageTwo",
    kwargs = dict(
        split="val",
    )
)

register(
    id = "msgr-test-v2",
    entry_point="messenger.envs:StageTwo",
    kwargs = dict(
        split="test",
    )
)

register(
    id = "msgr-test-se-v2",
    entry_point="messenger.envs:StageTwo",
    kwargs = dict(
        split="test_se",
    )
)

register(
    id = "msgr-train-v3",
    entry_point="messenger.envs:TwoEnvWrapper",
    kwargs = dict(
        stage=3,
        split_1="train_mc",
        split_2="train_sc",
        prob_env_1=0.75,
    )
)

register(
    id = "msgr-train-sc-v3",
    entry_point="messenger.envs:StageThree",
    kwargs = dict(
        split="train_sc",
    )
)

register(
    id = "msgr-train-mc-v3",
    entry_point="messenger.envs:StageThree",
    kwargs = dict(
        split="train_mc",
    )
)


register(
    id = "msgr-val-v3",
    entry_point="messenger.envs:StageThree",
    kwargs = dict(
        split="val",
    )
)

register(
    id = "msgr-test-v3",
    entry_point="messenger.envs:StageThree",
    kwargs = dict(
        split="test",
    )
)
src/messenger/__init__.py
microsoft/SmartPlay
from messenger.envs.wrappers import TwoEnvWrapper
from messenger.envs.stage_one import StageOne
from messenger.envs.stage_two import StageTwo
from messenger.envs.stage_three import StageThree
src/messenger/envs/__init__.py
microsoft/SmartPlay
import random
from collections import namedtuple

import gym
from gym import spaces
import numpy as np

import messenger.envs.config as config
from messenger.envs.config import Entity

# Positions of the entities
Position = namedtuple('Position', ["x", "y"])


class MessengerEnv(gym.Env):
    '''
    Base Messenger class that defines the action and observation spaces.
    '''

    def __init__(self):
        super().__init__()
        # up, down, left, right, stay
        self.action_space = spaces.Discrete(len(config.ACTIONS))

        # observations, not including the text manual
        self.observation_space = spaces.Dict({
            "entities": spaces.Box(
                low=0,
                high=14,
                shape=(config.STATE_HEIGHT, config.STATE_WIDTH, 3)
            ),
            "avatar": spaces.Box(
                low=15,
                high=16,
                shape=(config.STATE_HEIGHT, config.STATE_WIDTH, 1)
            )
        })

    def reset(self):
        raise NotImplementedError

    def step(self):
        raise NotImplementedError

    def render(self):
        raise NotImplementedError


class Grid:
    '''
    Class which makes it easier to build a grid observation from the dict state
    return by VGDLEnv.
    '''
    def __init__(self, layers, shuffle=True):
        '''
        layers:
            Each add() operation will place a separate entity in a new layer.
            Thus, this is the upper-limit to the number of items to be added.
        shuffle:
            Place each items in a random order.
        '''
        self.grid = np.zeros((config.STATE_HEIGHT, config.STATE_WIDTH, layers))
        self.order = list(range(layers)) # insertion order
        if shuffle:
            random.shuffle(self.order)
        self.layers = layers
        self.entity_count = 0

    def add(self, entity:Entity, position:Position):
        '''
        Add entity entity and position position.
        '''
        assert self.entity_count < self.layers, \
            f"Tried to add entity no. {self.entity_count} with {self.layers} layers."

        self.grid[position.y, position.x, self.order[self.entity_count]] = entity.id
        self.entity_count += 1
src/messenger/envs/base.py
microsoft/SmartPlay
import random
import json
from collections import namedtuple

import messenger.envs.config as config

Descr = namedtuple("Description", ['entity', 'role', 'type'])

class TextManual:
    '''
    Class which implements methods that allow environments to construct
    text manuals for games in Messenger.
    '''
    
    def __init__(self, json_path):
        with open(json_path, "r") as f:
            self.descriptors = json.load(f)

    def get_descriptor(self, entity, role, entity_type, no_type_p=0.15):
        '''
        Get a descriptor using the templates.
        Parameters:
        entity: The object that is being described (e.g. alien)
        role: The role the the object plays in the env. (e.g. enemy)
        entity_type:
            The object type (e.g. chaser)
        no_type_p:
            The probability of returning a descriptor that does not have
            any type information (only if entity_type is not None).
        '''
        if random.random() < no_type_p: # no type information
            return random.choice(self.descriptors[entity][role]["unknown"])

        else:
            return random.choice(self.descriptors[entity][role][entity_type])

    def get_document(self, enemy, message, goal, shuffle=False, 
                enemy_type=None, message_type=None, goal_type=None,
                append=False, delete=False, **kwargs):
        '''
        Makes a document for Messenger using the specified parameters.
        If no type is provided, a random type will be selected.

        Parameters:
        append: 
            If True, append an extraneous sentence to the document describing a 
            random object that is not in {enemy, message, goal}.
        delete: If True, Delete a random descriptor from the document.
        shuffle: 
            If True, shuffles the order of the descriptors
        kwargs:
            All other kwargs go to get_descriptor()
        '''

        document = [
            self.get_descriptor(entity=enemy, entity_type=enemy_type, role="enemy", **kwargs),
            self.get_descriptor(entity=message, entity_type=message_type, role="message", **kwargs),
            self.get_descriptor(entity=goal, entity_type=goal_type, role="goal", **kwargs)
        ]

        if delete: # delete a random descriptor
            document = random.sample(document, 2)

        if append:
            # choose an object not in {enemy, message, goal}
            valid_objs = [obj.name for obj in config.NPCS if obj.name not in [enemy, message, goal]]
            rand_obj = random.choice(valid_objs)
            result = None
            while result is None:
                try:
                    result = self.get_descriptor(
                        entity=rand_obj,
                        role=random.choice(("enemy", "message", "goal")),
                        entity_type=random.choice(("chaser", "fleeing", "immovable")),
                        **kwargs
                    )
                except:
                    pass
            document.append(result)
            
        if shuffle:
            document = random.sample(document, len(document))

        return document

    def get_document_plus(self, *args, **kwargs):
        '''
        Makes a document for Messenger using the specified parameters.

        Parameters:
        args: List of Descrip namedtuples
        '''

        document = []

        for descrip in args:
            document.append(
                self.get_descriptor(
                    entity=descrip.entity,
                    role=descrip.role,
                    entity_type=descrip.type,
                    no_type_p=0,
                    **kwargs
                )
            )

        return document

    def get_decoy_descriptor(self, entity, not_of_role, not_of_type, **kwargs):
        '''
        Get a description about the entity where the entity is not of role not_of_role
        and not of type not_of_type
        '''
        possible_roles = [x for x in ('message', 'goal', 'enemy') if x != not_of_role]
        random.shuffle(possible_roles)
        selected_type = random.choice([x for x in ('chaser', 'fleeing', 'immovable') if x != not_of_type])

        for role in possible_roles:
            try:
                return self.get_descriptor(
                    entity=entity,
                    role=role,
                    entity_type=selected_type,
                    no_type_p=0,
                    **kwargs
                )
            except:
                continue
        raise Exception('decoy description with impossible constraints')


if __name__ == "__main__":
    # just some quick and dirty tests
    from pathlib import Path
    text_json = Path(__file__).parent.joinpath('texts', 'text_train.json')
    manual = TextManual(json_path=text_json)
    descriptions = (
        Descr(entity="airplane", role='goal', type='chaser'),
        Descr(entity="airplane", role='message', type='fleeing'),
        Descr(entity="dog", role='enemy', type='immovable'),
        Descr(entity="dog", role='goal', type='fleeing'),
        Descr(entity="mage", role='goal', type='chaser'),
        Descr(entity="mage", role='message', type='immovable'),
    )
    print(manual.get_document_plus(*descriptions))
src/messenger/envs/manual.py
microsoft/SmartPlay
'''
Classes that follows a gym-like interface and implements stage one of the Messenger
environment.
'''

import json
import random
from collections import namedtuple
from pathlib import Path

import numpy as np

from messenger.envs.base import MessengerEnv, Position
import messenger.envs.config as config
from messenger.envs.manual import TextManual
from messenger.envs.utils import games_from_json


# Used to track sprites in StageOne, where we do not use VGDL to handle sprites.
Sprite = namedtuple("Sprite", ["name", "id", "position"])


class StageOne(MessengerEnv):
    def __init__(self, split, message_prob=0.2, shuffle_obs=True):
        '''
        Stage one where objects are all immovable. Since the episode length is short and entities
        do not move, we do not use VGDL engine for efficiency.
        message_prob:
            the probability that the avatar starts with the message
        shuffle_obs:
            shuffle the observation including the text manual
        '''
        super().__init__()
        self.message_prob = message_prob
        self.shuffle_obs = shuffle_obs
        this_folder = Path(__file__).parent
        
        # Get the games and manual
        games_json_path = this_folder.joinpath("games.json")
        if "train" in split and "mc" in split: # multi-combination games
            game_split = "train_multi_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "train" in split and "sc" in split: # single-combination games
            game_split = "train_single_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "val" in split:
            game_split = "val"
            text_json_path = this_folder.joinpath("texts", "text_val.json")
        elif "test" in split:
            game_split = "test"
            text_json_path = this_folder.joinpath("texts", "text_test.json")
        else:
            raise Exception(f"Split: {split} not understood.")

        # list of Game namedtuples
        self.all_games = games_from_json(json_path=games_json_path, split=game_split)
        
        # we only need the immovable and unknown descriptions, so just extract those.
        with text_json_path.open(mode="r") as f:
            descrip = json.load(f)
        
        self.descriptors = {}
        for entity in descrip:
            self.descriptors[entity] = {}
            for role in ("enemy", "message", "goal"):
                self.descriptors[entity][role] = []
                for sent in descrip[entity][role]["immovable"]:
                    self.descriptors[entity][role].append(sent)
                for sent in descrip[entity][role]["unknown"]:
                    self.descriptors[entity][role].append(sent)
        
        self.positions = [ # all possible entity locations
            Position(y=3, x=5),
            Position(y=5, x=3),
            Position(y=5, x=7),
            Position(y=7, x=5)
        ]
        self.avatar_start_pos = Position(y=5, x=5)
        self.avatar = None
        self.enemy = None
        self.message = None
        self.neutral = None
        self.goal = None

    def _get_manual(self):
        enemy_str = random.choice(self.descriptors[self.enemy.name]["enemy"])
        key_str = random.choice(self.descriptors[self.message.name]["message"])
        goal_str = random.choice(self.descriptors[self.goal.name]["goal"])
        manual = [enemy_str, key_str, goal_str]
        if self.shuffle_obs:
            random.shuffle(manual)
        return manual

    def _get_obs(self):
        entities = np.zeros((config.STATE_HEIGHT, config.STATE_WIDTH, 1))
        avatar = np.zeros((config.STATE_HEIGHT, config.STATE_WIDTH, 1))
        for sprite in (self.enemy, self.message, self.goal):
            entities[sprite.position.y, sprite.position.x, 0] = sprite.id
            
        avatar[self.avatar.position.y, self.avatar.position.x, 0] = self.avatar.id
        
        return {"entities": entities, "avatar": avatar}

    def reset(self):
        self.game = random.choice(self.all_games)
        enemy, message, goal = self.game.enemy, self.game.message, self.game.goal

        # randomly choose where to put enemy, key, goal
        shuffled_pos = random.sample(self.positions, 4)
        self.enemy = Sprite(name=enemy.name, id=enemy.id, position=shuffled_pos[0])
        self.message = Sprite(name=message.name, id=message.id, position=shuffled_pos[1])
        self.goal = Sprite(name=goal.name, id=goal.id, position=shuffled_pos[2])
        
        if random.random() < self.message_prob:
            self.avatar = Sprite(
                name=config.WITH_MESSAGE.name,
                id=config.WITH_MESSAGE.id,
                position=self.avatar_start_pos
            )

        else: # decide whether avatar has message or not
            self.avatar = Sprite(
                name=config.NO_MESSAGE.name,
                id=config.NO_MESSAGE.id,
                position=self.avatar_start_pos
            )
        
        obs = self._get_obs()
        manual = self._get_manual()

        return obs, manual
    
    def _move_avatar(self, action):
        if action == config.ACTIONS.stay:
            return
        
        elif action == config.ACTIONS.up: 
            if self.avatar.position.y <= 0:
                return
            else:
                new_position = Position(
                    y = self.avatar.position.y - 1,
                    x = self.avatar.position.x
                )
                
        elif action == config.ACTIONS.down: 
            if self.avatar.position.y >= config.STATE_HEIGHT - 1:
                return
            else:
                new_position = Position(
                    y = self.avatar.position.y + 1,
                    x = self.avatar.position.x
                )
                
        elif action == config.ACTIONS.left: 
            if self.avatar.position.x <= 0:
                return
            else:
                new_position = Position(
                    y = self.avatar.position.y,
                    x = self.avatar.position.x - 1
                )
                
        elif action == config.ACTIONS.right: 
            if self.avatar.position.x >= config.STATE_WIDTH - 1:
                return
            else:
                new_position = Position(
                    y = self.avatar.position.y,
                    x = self.avatar.position.x + 1
                )
                
        else:
            raise Exception(f"{action} is not a valid action.")
            
        self.avatar = Sprite(
                name=self.avatar.name,
                id=self.avatar.id,
                position=new_position
            )
            
    def _overlap(self, sprite_1, sprite_2):
        if (sprite_1.position.x == sprite_2.position.x and
           sprite_1.position.y == sprite_2.position.y):
            return True
        else:
            return False

    def step(self, action):
        self._move_avatar(action)
        obs = self._get_obs()
        if self._overlap(self.avatar, self.enemy):
            return obs, -1.0, True, None  # state, reward, done, info
        
        if self._overlap(self.avatar, self.message):
            if self.avatar.name == config.WITH_MESSAGE.name:
                return obs, -1.0, True, None
            elif self.avatar.name == config.NO_MESSAGE.name:
                return obs, 1.0, True, None
            else:
                raise Exception("Unknown avatar name {avatar.name}")
            
        if self._overlap(self.avatar, self.goal):
            if self.avatar.name == config.WITH_MESSAGE.name:
                return obs, 1.0, True, None
            elif self.avatar.name == config.NO_MESSAGE.name:
                return obs, -1.0, True, None
            else:
                raise Exception("Unknown avatar name {avatar.name}")
        
        return obs, 0.0, False, None
src/messenger/envs/stage_one.py
microsoft/SmartPlay
'''
Classes that follows a gym-like interface and implements stage three of the Messenger
environment.
'''

import json
import random
from collections import namedtuple
from pathlib import Path
from os import environ
import re

# hack to stop PyGame from printing to stdout
environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "hide"

import numpy as np
from vgdl.interfaces.gym import VGDLEnv

from messenger.envs.base import MessengerEnv, Grid, Position
import messenger.envs.config as config
from messenger.envs.manual import TextManual, Descr
from messenger.envs.utils import games_from_json


# specifies the game variant path is path to the vgdl domain file describing the variant.
GameVariant = namedtuple(
    "GameVariant",
    [
        "path",
        "enemy_type",
        "message_type",
        "goal_type",
        "decoy_message_type",
        "decoy_goal_type"
    ]
)

class StageThree(MessengerEnv):
    '''
    Similar to stage two Messenger, except with decoy objects that require
    disambiduation (e.g. chasing knight, vs immovable knight)
    '''

    def __init__(self, split:str, shuffle_obs=True):
        super().__init__()
        self.shuffle_obs = shuffle_obs # shuffle the entity layers

        this_folder = Path(__file__).parent
        # Get the games and manual
        games_json_path = this_folder.joinpath("games.json")
        if "train" in split and "mc" in split: # multi-combination games
            game_split = "train_multi_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "train" in split and "sc" in split: # single-combination games
            game_split = "train_single_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "val" in split:
            game_split = "val"
            text_json_path = this_folder.joinpath("texts", "text_val.json")
        elif "test" in split:
            game_split = "test"
            text_json_path = this_folder.joinpath("texts", "text_test.json")
        else:
            raise Exception(f"Split: {split} not understood.")

        # list of Game namedtuples
        self.all_games = games_from_json(json_path=games_json_path, split=game_split)
        self.text_manual = TextManual(json_path=text_json_path)

        vgdl_files = this_folder.joinpath("vgdl_files", "stage_3")
        
        # get the file paths to possible starting states
        self.init_states = [
            str(path) for path in vgdl_files.joinpath("init_states").glob("*.txt")
        ]
        # get all the game variants
        self.game_variants = [
            self._get_variant(path) for path in vgdl_files.joinpath("variants").glob("*.txt")
        ]

        # entities tracked by VGDLEnv
        self.notable_sprites = ["enemy", "message", "goal", "decoy_message", "decoy_goal", "no_message", "with_message"]
        self.env = None # the VGDLEnv

    def _get_variant(self, variant_file:Path) -> GameVariant:
        '''
        Return the GameVariant for the variant specified by variant_file. 
        Searches through the vgdl code to find the correct type:
        {chaser, fleeing, immovable}
        '''

        code = variant_file.read_text()
        return GameVariant(
            path = str(variant_file),
            enemy_type = re.search(r'enemy > (\S+)', code)[1].lower(),
            message_type = re.search(r'message > (\S+)', code)[1].lower(),
            goal_type = re.search(r'goal > (\S+)', code)[1].lower(),
            decoy_message_type = re.search(r'decoy_message > (\S+)', code)[1].lower(),
            decoy_goal_type = re.search(r'decoy_goal > (\S+)', code)[1].lower()
        )

    def _convert_obs(self, vgdl_obs):
        '''
        Return a grid built from the vgdl observation which is a
        KeyValueObservation object (see vgdl code for details).
        '''
        entity_locs = Grid(layers=5, shuffle=self.shuffle_obs)
        avatar_locs = Grid(layers=1)

        # try to add each entity one by one, if it's not there move on.
        if 'enemy.1' in vgdl_obs:
            entity_locs.add(self.game.enemy, Position(*vgdl_obs['enemy.1']['position']))
        if 'message.1' in vgdl_obs:
            entity_locs.add(self.game.message, Position(*vgdl_obs['message.1']['position']))
        else:
            # advance the entity counter, Oracle model requires special order.
            # TODO: maybe used named layers to make this more understandable.
            entity_locs.entity_count += 1
        if 'goal.1' in vgdl_obs:
            entity_locs.add(self.game.goal, Position(*vgdl_obs['goal.1']['position']))
        
        if 'decoy_message.1' in vgdl_obs:
            entity_locs.add(self.game.message, Position(*vgdl_obs['decoy_message.1']['position']))
        if 'decoy_goal.1' in vgdl_obs:
            entity_locs.add(self.game.goal, Position(*vgdl_obs['decoy_goal.1']['position']))

        if 'no_message.1' in vgdl_obs:
            '''
            Due to a quirk in VGDL, the avatar is no_message if it starts as no_message
            even if the avatar may have acquired the message at a later point.
            To check, if it has a message, check that the class vector corresponding to
            with_message is == 1.
            '''
            avatar_pos = Position(*vgdl_obs['no_message.1']['position'])
            # with_key is last position, see self.notable_sprites
            if vgdl_obs['no_message.1']['class'][-1] == 1:
                avatar = config.WITH_MESSAGE
            else:
                avatar = config.NO_MESSAGE

        elif "with_message.1" in vgdl_obs:
            # this case only occurs if avatar begins as with_message at start of episode
            avatar_pos = Position(*vgdl_obs['with_message.1']['position'])
            avatar = config.WITH_MESSAGE

        else: # the avatar is not in observation, so is probably dead
            return {"entities": entity_locs.grid, "avatar": avatar_locs.grid}

        avatar_locs.add(avatar, avatar_pos) # if not dead, add it.

        return {"entities": entity_locs.grid, "avatar": avatar_locs.grid}


    def reset(self, variant_id:int=None, **kwargs):
        '''
        Resets the current environment. NOTE: We remake the environment each time.
        This is a workaround to a bug in py-vgdl, where env.reset() does not
        properly reset the environment. kwargs go to get_document().
        '''

        self.game = random.choice(self.all_games) # (e.g. enemy-alien, message-knight, goal - bear)

        # choose the game variant (e.g. enmey-chasing, message-fleeing, goal-static)
        # and initial starting location of the entities.
        if variant_id is not None:
            variant = self.game_variants[variant_id]
        else:
            variant = random.choice(self.game_variants)
        init_state = random.choice(self.init_states) # inital state file

        # args that will go into VGDL Env.
        self._envargs = {
            'game_file': variant.path,
            'level_file': init_state,
            'notable_sprites': self.notable_sprites.copy(),
            'obs_type': 'objects', # track the objects
            'block_size': 34  # rendering block size
        }
        self.env = VGDLEnv(**self._envargs)
        vgdl_obs = self.env.reset()

        all_npcs = (
            Descr(entity=self.game.enemy.name, role='enemy', type=variant.enemy_type),
            Descr(entity=self.game.message.name, role='message', type=variant.message_type),
            Descr(entity=self.game.goal.name, role='goal', type=variant.goal_type),
            Descr(entity=self.game.message.name, role='enemy', type=variant.decoy_message_type),
            Descr(entity=self.game.goal.name, role='enemy', type=variant.decoy_goal_type),
        )

        manual = self.text_manual.get_document_plus(*all_npcs, **kwargs)
        manual.append(
            self.text_manual.get_decoy_descriptor(
                entity=self.game.enemy.name,
                not_of_role="enemy",
                not_of_type=variant.enemy_type
            )
        )

        if self.shuffle_obs:
            random.shuffle(manual)

        return self._convert_obs(vgdl_obs), manual

    def step(self, action):
        vgdl_obs, reward, done, info = self.env.step(action)
        return self._convert_obs(vgdl_obs), reward, done, info
src/messenger/envs/stage_three.py
microsoft/SmartPlay
'''
Classes that follows a gym-like interface and implements stage two of the Messenger
environment.
'''

import json
import random
from collections import namedtuple
from pathlib import Path
from os import environ
import re

# hack to stop PyGame from printing to stdout
environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "hide"

from vgdl.interfaces.gym import VGDLEnv
import numpy as np

from messenger.envs.base import MessengerEnv, Grid, Position
import messenger.envs.config as config
from messenger.envs.manual import TextManual
from messenger.envs.utils import games_from_json


# specifies the game variant (e.g. chasing enemy, fleeing message, stationary goal)
# path is path to the vgdl domain file describing the variant.
GameVariant = namedtuple(
    "GameVariant", ["path", "enemy_type", "message_type", "goal_type"]
)


class StageTwo(MessengerEnv):
    '''
    Full messenger environment with mobile sprites. Uses Py-VGDL as game engine.
    To avoid the need to instantiate a large number of games, (since there are
    P(12,3) = 1320 possible entity to role assignments) We apply a wrapper on top
    of the text and game state which masks the role archetypes (enemy, message goal)
    into entities (e.g. alien, knight, mage).
    '''

    def __init__(self, split:str, shuffle_obs=True):
        super().__init__()
        self.shuffle_obs = shuffle_obs # shuffle the entity layers

        this_folder = Path(__file__).parent
        # Get the games and manual
        games_json_path = this_folder.joinpath("games.json")
        if "train" in split and "mc" in split: # multi-combination games
            game_split = "train_multi_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "train" in split and "sc" in split: # single-combination games
            game_split = "train_single_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "val" in split:
            game_split = "val"
            text_json_path = this_folder.joinpath("texts", "text_val.json")
        elif "test" in split:
            game_split = "test"
            text_json_path = this_folder.joinpath("texts", "text_test.json")
        else:
            raise Exception(f"Split: {split} not understood.")

        # list of Game namedtuples
        self.all_games = games_from_json(json_path=games_json_path, split=game_split)
        self.text_manual = TextManual(json_path=text_json_path)

        # get the folder that has the game variants and init_states
        if "test" in split and "se" not in split: # new dynamics (se for state estimation)
            vgdl_files = this_folder.joinpath("vgdl_files", "stage_2_nd")
        else: # training dynamics
            vgdl_files = this_folder.joinpath("vgdl_files", "stage_2")
        
        # get the file paths to possible starting states
        self.init_states = [
            str(path) for path in vgdl_files.joinpath("init_states").glob("*.txt")
        ]
        # get all the game variants
        self.game_variants = [
            self._get_variant(path) for path in vgdl_files.joinpath("variants").glob("*.txt")
        ]

        # entities tracked by VGDLEnv
        self.notable_sprites = ["enemy", "message", "goal", "no_message", "with_message"]
        self.env = None # the VGDLEnv

    def _get_variant(self, variant_file:Path) -> GameVariant:
        '''
        Return the GameVariant for the variant specified by variant_file. 
        Searches through the vgdl code to find the correct type:
        {chaser, fleeing, immovable}
        '''

        code = variant_file.read_text()
        return GameVariant(
            path = str(variant_file),
            enemy_type = re.search(r'enemy > (\S+)', code)[1].lower(),
            message_type = re.search(r'message > (\S+)', code)[1].lower(),
            goal_type = re.search(r'goal > (\S+)', code)[1].lower()
        )

    def _convert_obs(self, vgdl_obs):
        '''
        Return a grid built from the vgdl observation which is a
        KeyValueObservation object (see vgdl code for details).
        '''
        entity_locs = Grid(layers=3, shuffle=self.shuffle_obs)
        avatar_locs = Grid(layers=1)

        # try to add each entity one by one, if it's not there move on.
        if 'enemy.1' in vgdl_obs:
            entity_locs.add(self.game.enemy, Position(*vgdl_obs['enemy.1']['position']))
        if 'message.1' in vgdl_obs:
            entity_locs.add(self.game.message, Position(*vgdl_obs['message.1']['position']))
        else:
            # advance the entity counter, Oracle model requires special order.
            # TODO: maybe used named layers to make this more understandable.
            entity_locs.entity_count += 1
        if 'goal.1' in vgdl_obs:
            entity_locs.add(self.game.goal, Position(*vgdl_obs['goal.1']['position']))

        if 'no_message.1' in vgdl_obs:
            '''
            Due to a quirk in VGDL, the avatar is no_message if it starts as no_message
            even if the avatar may have acquired the message at a later point.
            To check, if it has a message, check that the class vector corresponding to
            with_message is == 1.
            '''
            avatar_pos = Position(*vgdl_obs['no_message.1']['position'])
            # with_key is last position, see self.notable_sprites
            if vgdl_obs['no_message.1']['class'][-1] == 1:
                avatar = config.WITH_MESSAGE
            else:
                avatar = config.NO_MESSAGE

        elif "with_message.1" in vgdl_obs:
            # this case only occurs if avatar begins as with_message at start of episode
            avatar_pos = Position(*vgdl_obs['with_message.1']['position'])
            avatar = config.WITH_MESSAGE

        else: # the avatar is not in observation, so is probably dead
            return {"entities": entity_locs.grid, "avatar": avatar_locs.grid}

        avatar_locs.add(avatar, avatar_pos) # if not dead, add it.

        return {"entities": entity_locs.grid, "avatar": avatar_locs.grid}

    def reset(self, **kwargs):
        '''
        Resets the current environment. NOTE: We remake the environment each time.
        This is a workaround to a bug in py-vgdl, where env.reset() does not
        properly reset the environment. kwargs go to get_document().
        '''

        self.game = random.choice(self.all_games) # (e.g. enemy-alien, message-knight, goal - bear)

        # choose the game variant (e.g. enmey-chasing, message-fleeing, goal-static)
        # and initial starting location of the entities.
        variant = random.choice(self.game_variants)
        init_state = random.choice(self.init_states) # initial state file

        # args that will go into VGDL Env.
        self._envargs = {
            'game_file': variant.path,
            'level_file': init_state,
            'notable_sprites': self.notable_sprites.copy(),
            'obs_type': 'objects', # track the objects
            'block_size': 34  # rendering block size
        }
        self.env = VGDLEnv(**self._envargs)
        vgdl_obs = self.env.reset()

        manual = self.text_manual.get_document(
            enemy=self.game.enemy.name,
            message=self.game.message.name,
            goal=self.game.goal.name,
            enemy_type=variant.enemy_type,
            message_type=variant.message_type,
            goal_type=variant.goal_type,
            **kwargs
        )

        if self.shuffle_obs:
            random.shuffle(manual)
            
        return self._convert_obs(vgdl_obs), manual

    def step(self, action):
        vgdl_obs, reward, done, info = self.env.step(action)
        return self._convert_obs(vgdl_obs), reward, done, info
    

src/messenger/envs/stage_two.py
microsoft/SmartPlay
'''
Implements wrappers on top of the basic messenger environments
'''
import random

from messenger.envs.base import MessengerEnv
from messenger.envs.stage_one import StageOne
from messenger.envs.stage_two import StageTwo
from messenger.envs.stage_three import StageThree


class TwoEnvWrapper(MessengerEnv):
    '''
    Switches between two Messenger environments
    '''
    def __init__(self, stage:int, split_1:str, split_2:str, prob_env_1=0.5, **kwargs):
        super().__init__()
        if stage == 1:
            self.env_1 = StageOne(split=split_1, **kwargs)
            self.env_2 = StageOne(split=split_2, **kwargs)
        elif stage == 2:
            self.env_1 = StageTwo(split=split_1, **kwargs)
            self.env_2 = StageTwo(split=split_2, **kwargs)
        elif stage == 3:
            self.env_1 = StageThree(split=split_1, **kwargs)
            self.env_2 = StageThree(split=split_2, **kwargs)
        
        self.prob_env_1 = prob_env_1
        self.cur_env = None
        
    def reset(self, **kwargs):
        if random.random() < self.prob_env_1:
            self.cur_env = self.env_1
        else:
            self.cur_env = self.env_2
        return self.cur_env.reset(**kwargs)
    
    def step(self, action):
        return self.cur_env.step(action)
src/messenger/envs/wrappers.py
microsoft/SmartPlay
'''
Common code and utilities used by the models
'''

import torch

class ObservationBuffer:
    '''
    Maintains a buffer of observations along the 0-dim. Observations
    are currently expected to be a dict of np arrays. Currently keeps
    observations in a list and then stacks them via torch.stack().
    TODO: pre-allocate memory for faster calls to get_obs().
    
    Parameters:
    buffer_size
        How many previous observations to track in the buffer
    device
        The device on which buffers are loaded into
    '''
    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.buffer = None
        self.device = device

    def _np_to_tensor(self, obs):
        return torch.from_numpy(obs).long().to(self.device)

    def reset(self, obs):
        # initialize / reset the buffer with the observation
        self.buffer = [obs for _ in range(self.buffer_size)]

    def update(self, obs):
        # update the buffer by appending newest observation
        assert self.buffer, "Please initialize buffer first with reset()"
        del self.buffer[0] # delete the oldest entry
        self.buffer.append(obs) # append the newest observation

    def get_obs(self):
        # get a stack of all observations currently in the buffer
        stacked_obs = {}
        for key in self.buffer[0].keys():
            stacked_obs[key] = torch.stack(
                [self._np_to_tensor(obs[key]) for obs in self.buffer]
            )
        return stacked_obs


class Encoder:
    '''
    Text-encoder class with caching for fast sentence-encoding.
    self.encoder and self.tokenizer are expected to be HuggingFace model
    and its respective tokenizer. 
    
    Warning: currently have not implemented max cache size, watch out for
    out of memory errors if the number of possible inputs is v. large. All 
    original Messenger descriptions combined should take < 2GB of GPU memory.
    '''
    def __init__(self, model, tokenizer, device: torch.device, max_length:int=36):
        self.encoder = model.to(device)
        self.tokenizer = tokenizer
        self.device = device
        self.max_length = max_length # max sentence length
        self.cache = {}

    def to(self, device):
        self.device = device
        self.encoder = self.encoder.to(device)
        return self

    def tokens_to_device(self, tokens):
        tok_device = {}
        for key in tokens:
            tok_device[key] = tokens[key].to(self.device)
        return tok_device

    def encode(self, text):
        '''
        Encodes the text using self.encoder and self.tokenizer. Text should be
        a list of sents, where sent is a string.
        '''
        encoded = [] # the final encoded texts
        for sent in text:
            if sent in self.cache.keys(): # sentence is in cache
                encoded.append(self.cache[sent])
            else: 
                with torch.no_grad():
                    tokens = self.tokenizer(
                        sent,
                        return_tensors="pt",
                        truncation=False,
                        truncation_strategy='do_not_truncate',
                        padding="max_length",
                        max_length=self.max_length
                    )
                    emb = self.encoder(**self.tokens_to_device(tokens)).last_hidden_state
                encoded.append(emb)
                self.cache[sent] = emb
        return torch.cat(encoded, dim=0)

def nonzero_mean(emb):
    '''
    Takes as input an embedding, emb. It should be H x W x L x D. with
    optional batch dimension. (H,W) is the grid dim, L the layers and
    D the embedding dimension. Returns mean of non-zero vectors along L dim.
    This is used to take care of overlapping sprites.
    '''
    # Count the number of non-zero vectors
    non_zero = torch.sum(torch.norm(emb, dim=-1) > 0, dim=-1)
    non_zero = non_zero.unsqueeze(-1).float()  # broadcasting
    non_zero[non_zero == 0] = 1 # prevent division by zero
    return torch.sum(emb, dim=-2) / non_zero
src/messenger/models/utils.py
microsoft/SmartPlay
# This file is used to load all games in the smartplay directory.

import warnings
import importlib
import os
import yaml


root_dir = os.path.dirname(os.path.abspath(__file__))
games = []
env_list = []


_exclude_path = ['__pycache__', 'utils', 'tests', 'eval', 'example_game']
_module_dir = os.path.dirname(__file__)


for dirname in os.listdir(root_dir):
    if os.path.isdir(os.path.join(root_dir, dirname)) and dirname not in _exclude_path:
        if '__init__.py' in os.listdir(os.path.join(root_dir, dirname)):
            games.append(dirname)


game_challenges = {}
recorded_settings = {}


for game in games:


    if not os.path.exists(os.path.join(root_dir, game, 'evaluation.yml')):
        warnings.warn('Game `{}` does not have evaluation.yml. Skipping the game.'.format(game), UserWarning)
        continue


    try:
        # Load evaluation settings
        with open(os.path.join(root_dir, game, 'evaluation.yml'), 'r') as f:
            yml = yaml.safe_load(f)
            recorded_setting = yml['recorded settings']
            game_challenge = yml['challenges']

    except Exception as e:
        warnings.warn('Failed to load `evaluation.yml` for `{}`.'.format(game), UserWarning)
        continue


    try:
        # Load environments
        module = importlib.import_module("."+game, package='smartplay')
        environments = getattr(module, 'environments', None)
        if environments:
            for env_name, version in environments:
                env_list.append('{}-{}'.format(env_name, version))
                game_challenges[env_name] = game_challenge['all'] if env_name not in game_challenge.keys() else game_challenge[env_name]
                if env_name in recorded_setting.keys():
                    recorded_settings[env_name] = recorded_setting[env_name]
        else:
            warnings.warn('Failed to load `{}.environments`. Skipping the game.'.format(game), UserWarning)
            continue

    except Exception as e:
        warnings.warn('Failed to import `{}`. Skipping the game.'.format(game), UserWarning)
        continue


from .eval import *
src/smartplay/__init__.py
microsoft/SmartPlay
from gym.envs.registration import register

from .bandit import BanditTwoArmedDeterministicFixed
from .bandit import BanditTwoArmedHighHighFixed
from .bandit import BanditTwoArmedHighLowFixed
from .bandit import BanditTwoArmedLowLowFixed

environments = [['BanditTwoArmedDeterministicFixed', 'v0'],
                ['BanditTwoArmedHighHighFixed', 'v0'],
                ['BanditTwoArmedHighLowFixed', 'v0'],
                ['BanditTwoArmedLowLowFixed', 'v0']]

for environment in environments:
    register(
        id='{}-{}'.format(environment[0], environment[1]),
        entry_point='smartplay.bandits:{}'.format(environment[0]),
    )

src/smartplay/bandits/__init__.py
microsoft/SmartPlay
import numpy as np
import gym
from gym import spaces
from gym.utils import seeding
from ..utils import HistoryTracker, describe_act
import random

class BanditEnv(gym.Env):
    """
    Bandit environment base to allow agents to interact with the class n-armed bandit
    in different variations

    p_dist:
        A list of probabilities of the likelihood that a particular bandit will pay out
    r_dist:
        A list of either rewards (if number) or means and standard deviations (if list)
        of the payout that bandit has
    """

    default_iter = 20
    default_steps = 50

    def __init__(self, p_dist, r_dist, max_steps=50):
        if len(p_dist) != len(r_dist):
            raise ValueError("Probability and Reward distribution must be the same length")

        if min(p_dist) < 0 or max(p_dist) > 1:
            raise ValueError("All probabilities must be between 0 and 1")

        for reward in r_dist:
            if isinstance(reward, list) and reward[1] <= 0:
                raise ValueError("Standard deviation in rewards must all be greater than 0")

        self.n_bandits = len(p_dist)
        self.action_space = spaces.Discrete(self.n_bandits)
        self.action_list = ["Pull slot machine {}.".format(i+1) for i in range(self.n_bandits)]
        self.observation_space = spaces.Discrete(1)
        self.p_dist = p_dist
        self.r_dist = r_dist
        self.desc = """
You are in the casino with 2 slot machines in front of you. Your goal is to try to earn the most from those slot machines.

{}
""".format(describe_act(self.action_list)).strip()
        self.history = HistoryTracker(max_steps)

        self._seed()
        self.score_tracker = 0

    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, action):
        assert self.action_space.contains(action)

        reward = -1
        done = False

        if np.random.uniform() < self.p_dist[action]:
            if not isinstance(self.r_dist[action], list):
                reward = self.r_dist[action]
            else:
                reward = np.random.normal(self.r_dist[action][0], self.r_dist[action][1])

        self.score_tracker+=int(action == self.optimal)
        info = {"obs": "You pulled slot machine {}, you received reward {}.".format(action+1, reward),
                "score": self.score_tracker,
                "manual": self.desc,
                "history": self.history.describe(),
                "optimal":self.optimal,
                "completed": 0,
                }
        self.history.step(info)

        return 0, reward, done, info

    def reset(self):
        idx_list = list(range(len(self.p_dist)))
        random.shuffle(idx_list)

        self.score_tracker=0
        self.p_dist = [self.p_dist[i] for i in idx_list]
        self.r_dist = [self.r_dist[i] for i in idx_list]
        self.ev = [p*r - (1-p) for p, r in zip(self.p_dist, self.r_dist)]
        self.optimal = np.argmax(self.ev)
        self.history.reset()
        info = {"obs": "A new round begins.",
                "score": self.score_tracker,
                "manual": self.desc,
                "history": self.history.describe(),
                "optimal":self.optimal,
                "completed": 0,
                }
        self.history.step(info)
        return 0, info

    def render(self, mode='human', close=False):
        pass


class BanditTwoArmedDeterministicFixed(BanditEnv):
    """Simplest case where one bandit always pays, and the other always doesn't"""
    def __init__(self):
        BanditEnv.__init__(self, p_dist=[1, 0], r_dist=[1, 1])


class BanditTwoArmedHighLowFixed(BanditEnv):
    """Stochastic version with a large difference between which bandit pays out of two choices"""
    def __init__(self):
        BanditEnv.__init__(self, p_dist=[0.8, 0.2], r_dist=[1, 1])


class BanditTwoArmedHighHighFixed(BanditEnv):
    """Stochastic version with a small difference between which bandit pays where both are good"""
    def __init__(self):
        BanditEnv.__init__(self, p_dist=[0.8, 0.9], r_dist=[1, 1])


class BanditTwoArmedLowLowFixed(BanditEnv):
    """Stochastic version with a small difference between which bandit pays where both are bad"""
    def __init__(self):
        BanditEnv.__init__(self, p_dist=[0.1, 0.2], r_dist=[1, 1])

src/smartplay/bandits/bandit.py
microsoft/SmartPlay
from gym.envs.registration import register

from .crafter_env import Crafter

environments = [
    ['Crafter', 'v0'],
]

for environment in environments:
    register(
        id='{}-{}'.format(environment[0], environment[1]),
        entry_point='smartplay.crafter:{}'.format(environment[0]),
        kwargs={'reward': True}
    )

src/smartplay/crafter/__init__.py
microsoft/SmartPlay
from .env import Env, EnvSample
from .recorder import Recorder

# try:
#   import gym
#   gym.register(
#       id='CrafterReward-v1',
#       entry_point='crafter:Env',
#       max_episode_steps=10000,
#       kwargs={'reward': True})
#   gym.register(
#       id='CrafterNoReward-v1',
#       entry_point='crafter:Env',
#       max_episode_steps=10000,
#       kwargs={'reward': False})
# except ImportError:
#   pass

src/smartplay/crafter/crafter/__init__.py
microsoft/SmartPlay
import collections

import numpy as np

from . import constants
from . import engine
from . import objects
from . import worldgen


# Gym is an optional dependency.
try:
  import gym
  DiscreteSpace = gym.spaces.Discrete
  BoxSpace = gym.spaces.Box
  DictSpace = gym.spaces.Dict
  BaseClass = gym.Env
except ImportError:
  DiscreteSpace = collections.namedtuple('DiscreteSpace', 'n')
  BoxSpace = collections.namedtuple('BoxSpace', 'low, high, shape, dtype')
  DictSpace = collections.namedtuple('DictSpace', 'spaces')
  BaseClass = object


class Env(BaseClass):

  def __init__(
      self, area=(64, 64), view=(9, 9), size=(64, 64),
      reward=True, length=10000, seed=None):
    view = np.array(view if hasattr(view, '__len__') else (view, view))
    size = np.array(size if hasattr(size, '__len__') else (size, size))
    seed = np.random.randint(0, 2**31 - 1) if seed is None else seed
    self._area = area
    self._view = view
    self._size = size
    self._reward = reward
    self._length = length
    self._seed = seed
    self._episode = 0
    self._world = engine.World(area, constants.materials, (12, 12))
    self._textures = engine.Textures(constants.root / 'assets')
    item_rows = int(np.ceil(len(constants.items) / view[0]))
    self._local_view = engine.LocalView(
        self._world, self._textures, [view[0], view[1] - item_rows])
    self._item_view = engine.ItemView(
        self._textures, [view[0], item_rows])
    self._sem_view = engine.SemanticView(self._world, [
        objects.Player, objects.Cow, objects.Zombie,
        objects.Skeleton, objects.Arrow, objects.Plant])
    self._step = None
    self._player = None
    self._last_health = None
    self._unlocked = None
    # Some libraries expect these attributes to be set.
    self.reward_range = None
    self.metadata = None

  @property
  def observation_space(self):
    return BoxSpace(0, 255, tuple(self._size) + (3,), np.uint8)

  @property
  def action_space(self):
    return DiscreteSpace(len(constants.actions))

  @property
  def action_names(self):
    return constants.actions

  def reset(self):
    center = (self._world.area[0] // 2, self._world.area[1] // 2)
    self._episode += 1
    self._step = 0
    self._world.reset(seed=hash((self._seed, self._episode)) % (2 ** 31 - 1))
    self._update_time()
    self._player = objects.Player(self._world, center)
    self._last_health = self._player.health
    self._world.add(self._player)
    self._unlocked = set()
    worldgen.generate_world(self._world, self._player)
    return self._obs()

  def step(self, action):
    self._step += 1
    self._update_time()
    self._player.action = constants.actions[action]
    for obj in self._world.objects:
      if self._player.distance(obj) < 2 * max(self._view):
        obj.update()
    if self._step % 10 == 0:
      for chunk, objs in self._world.chunks.items():
        # xmin, xmax, ymin, ymax = chunk
        # center = (xmax - xmin) // 2, (ymax - ymin) // 2
        # if self._player.distance(center) < 4 * max(self._view):
        self._balance_chunk(chunk, objs)
    obs = self._obs()
    reward = (self._player.health - self._last_health) / 10
    self._last_health = self._player.health
    unlocked = {
        name for name, count in self._player.achievements.items()
        if count > 0 and name not in self._unlocked}
    if unlocked:
      self._unlocked |= unlocked
      reward += 1.0
    dead = self._player.health <= 0
    over = self._length and self._step >= self._length
    done = dead or over
    info = {
        'inventory': self._player.inventory.copy(),
        'achievements': self._player.achievements.copy(),
        'sleeping': self._player.sleeping,
        'discount': 1 - float(dead),
        'semantic': self._sem_view(),
        'player_pos': self._player.pos,
        'player_facing': self._player.facing,
        'reward': reward,
        'dead': dead,
        'unlocked': unlocked,
        'action': self._player.action,
        'view': self._view,
    }
    if not self._reward:
      reward = 0.0
    return obs, reward, done, info

  def render(self, size=None):
    size = size or self._size
    unit = size // self._view
    canvas = np.zeros(tuple(size) + (3,), np.uint8)
    local_view = self._local_view(self._player, unit)
    item_view = self._item_view(self._player.inventory, unit)
    view = np.concatenate([local_view, item_view], 1)
    border = (size - (size // self._view) * self._view) // 2
    (x, y), (w, h) = border, view.shape[:2]
    canvas[x: x + w, y: y + h] = view
    return canvas.transpose((1, 0, 2))

  def _obs(self):
    return self.render()

  def _update_time(self):
    # https://www.desmos.com/calculator/grfbc6rs3h
    progress = (self._step / 300) % 1 + 0.3
    daylight = 1 - np.abs(np.cos(np.pi * progress)) ** 3
    self._world.daylight = daylight

  def _balance_chunk(self, chunk, objs):
    light = self._world.daylight
    self._balance_object(
        chunk, objs, objects.Zombie, 'grass', 6, 0, 0.3, 0.4,
        lambda pos: objects.Zombie(self._world, pos, self._player),
        lambda num, space: (
            0 if space < 50 else 3.5 - 3 * light, 3.5 - 3 * light))
    self._balance_object(
        chunk, objs, objects.Skeleton, 'path', 7, 7, 0.1, 0.1,
        lambda pos: objects.Skeleton(self._world, pos, self._player),
        lambda num, space: (0 if space < 6 else 1, 2))
    self._balance_object(
        chunk, objs, objects.Cow, 'grass', 5, 5, 0.01, 0.1,
        lambda pos: objects.Cow(self._world, pos),
        lambda num, space: (0 if space < 30 else 1, 1.5 + light))

  def _balance_object(
      self, chunk, objs, cls, material, span_dist, despan_dist,
      spawn_prob, despawn_prob, ctor, target_fn):
    xmin, xmax, ymin, ymax = chunk
    random = self._world.random
    creatures = [obj for obj in objs if isinstance(obj, cls)]
    mask = self._world.mask(*chunk, material)
    target_min, target_max = target_fn(len(creatures), mask.sum())
    if len(creatures) < int(target_min) and random.uniform() < spawn_prob:
      xs = np.tile(np.arange(xmin, xmax)[:, None], [1, ymax - ymin])
      ys = np.tile(np.arange(ymin, ymax)[None, :], [xmax - xmin, 1])
      xs, ys = xs[mask], ys[mask]
      i = random.randint(0, len(xs))
      pos = np.array((xs[i], ys[i]))
      empty = self._world[pos][1] is None
      away = self._player.distance(pos) >= span_dist
      if empty and away:
        self._world.add(ctor(pos))
    elif len(creatures) > int(target_max) and random.uniform() < despawn_prob:
      obj = creatures[random.randint(0, len(creatures))]
      away = self._player.distance(obj.pos) >= despan_dist
      if away:
        self._world.remove(obj)

class EnvSample(BaseClass):

  def __init__(
      self, area=(64, 64), view=(9, 9), size=(64, 64),
      reward=True, length=10000, seed=None, steps=20):
    view = np.array(view if hasattr(view, '__len__') else (view, view))
    size = np.array(size if hasattr(size, '__len__') else (size, size))
    seed = np.random.randint(0, 2**31 - 1) if seed is None else seed
    self._area = area
    self._view = view
    self._size = size
    self._reward = reward
    self._length = length
    self._seed = seed
    self._episode = 0
    self._world = engine.World(area, constants.materials, (12, 12))
    self._textures = engine.Textures(constants.root / 'assets')
    item_rows = int(np.ceil(len(constants.items) / view[0]))
    self._local_view = engine.LocalView(
        self._world, self._textures, [view[0], view[1] - item_rows])
    self._item_view = engine.ItemView(
        self._textures, [view[0], item_rows])
    self._sem_view = engine.SemanticView(self._world, [
        objects.Player, objects.Cow, objects.Zombie,
        objects.Skeleton, objects.Arrow, objects.Plant])
    self._step = None
    self._nsteps = steps
    self._player = None
    self._last_health = None
    self._unlocked = None
    # Some libraries expect these attributes to be set.
    self.reward_range = None
    self.metadata = None

  @property
  def observation_space(self):
    return BoxSpace(0, 255, tuple(self._size) + (3,), np.uint8)

  @property
  def action_space(self):
    return DiscreteSpace(len(constants.actions))

  @property
  def action_names(self):
    return constants.actions

  def reset(self):
    center = (self._world.area[0] // 2, self._world.area[1] // 2)
    self._episode += 1
    self._step = 0
    self._world.reset(seed=hash((self._seed, self._episode)) % (2 ** 31 - 1))
    self._update_time()
    self._player = objects.Player(self._world, center)
    self._last_health = self._player.health
    self._world.add(self._player)
    self._unlocked = set()
    worldgen.generate_world(self._world, self._player)
    return self._obs()

  def __step(self, action):
    self._step += 1
    self._update_time()
    self._player.action = constants.actions[action]
    for obj in self._world.objects:
      if self._player.distance(obj) < 2 * max(self._view):
        obj.update()
    if self._step % 10 == 0:
      for chunk, objs in self._world.chunks.items():
        # xmin, xmax, ymin, ymax = chunk
        # center = (xmax - xmin) // 2, (ymax - ymin) // 2
        # if self._player.distance(center) < 4 * max(self._view):
        self._balance_chunk(chunk, objs)
    obs = self._obs()
    reward = (self._player.health - self._last_health) / 10
    self._last_health = self._player.health
    unlocked = {
        name for name, count in self._player.achievements.items()
        if count > 0 and name not in self._unlocked}
    if unlocked:
      self._unlocked |= unlocked
      reward += 1.0
    dead = self._player.health <= 0
    over = self._length and self._step >= self._length
    done = dead or over
    info = {
        'inventory': self._player.inventory.copy(),
        'achievements': self._player.achievements.copy(),
        'discount': 1 - float(dead),
        'semantic': self._sem_view(),
        'player_pos': self._player.pos,
        'player_facing': self._player.facing,
        'reward': reward,
        'dead': dead,
        'unlocked': unlocked,
        'action': self._player.action,
        'view': self._view,
    }
    if not self._reward:
      reward = 0.0
    return obs, reward, done, info

  def step(self, action):
    frames = max(int(np.random.normal(loc = self._nsteps, scale = self._nsteps/2)), 1)
    for _ in range(frames):
      obs, reward, done, info = self.__step(np.random.randint(len(constants.actions)))
      if done:
        self.reset()
        obs, reward, done, info = self.__step(np.random.randint(len(constants.actions)))
    return obs, reward, done, info

  def render(self, size=None):
    size = size or self._size
    unit = size // self._view
    canvas = np.zeros(tuple(size) + (3,), np.uint8)
    local_view = self._local_view(self._player, unit)
    item_view = self._item_view(self._player.inventory, unit)
    view = np.concatenate([local_view, item_view], 1)
    border = (size - (size // self._view) * self._view) // 2
    (x, y), (w, h) = border, view.shape[:2]
    canvas[x: x + w, y: y + h] = view
    return canvas.transpose((1, 0, 2))

  def _obs(self):
    return self.render()

  def _update_time(self):
    # https://www.desmos.com/calculator/grfbc6rs3h
    progress = (self._step / 300) % 1 + 0.3
    daylight = 1 - np.abs(np.cos(np.pi * progress)) ** 3
    self._world.daylight = daylight

  def _balance_chunk(self, chunk, objs):
    light = self._world.daylight
    self._balance_object(
        chunk, objs, objects.Zombie, 'grass', 6, 0, 0.3, 0.4,
        lambda pos: objects.Zombie(self._world, pos, self._player),
        lambda num, space: (
            0 if space < 50 else 3.5 - 3 * light, 3.5 - 3 * light))
    self._balance_object(
        chunk, objs, objects.Skeleton, 'path', 7, 7, 0.1, 0.1,
        lambda pos: objects.Skeleton(self._world, pos, self._player),
        lambda num, space: (0 if space < 6 else 1, 2))
    self._balance_object(
        chunk, objs, objects.Cow, 'grass', 5, 5, 0.01, 0.1,
        lambda pos: objects.Cow(self._world, pos),
        lambda num, space: (0 if space < 30 else 1, 1.5 + light))

  def _balance_object(
      self, chunk, objs, cls, material, span_dist, despan_dist,
      spawn_prob, despawn_prob, ctor, target_fn):
    xmin, xmax, ymin, ymax = chunk
    random = self._world.random
    creatures = [obj for obj in objs if isinstance(obj, cls)]
    mask = self._world.mask(*chunk, material)
    target_min, target_max = target_fn(len(creatures), mask.sum())
    if len(creatures) < int(target_min) and random.uniform() < spawn_prob:
      xs = np.tile(np.arange(xmin, xmax)[:, None], [1, ymax - ymin])
      ys = np.tile(np.arange(ymin, ymax)[None, :], [xmax - xmin, 1])
      xs, ys = xs[mask], ys[mask]
      i = random.randint(0, len(xs))
      pos = np.array((xs[i], ys[i]))
      empty = self._world[pos][1] is None
      away = self._player.distance(pos) >= span_dist
      if empty and away:
        self._world.add(ctor(pos))
    elif len(creatures) > int(target_max) and random.uniform() < despawn_prob:
      obj = creatures[random.randint(0, len(creatures))]
      away = self._player.distance(obj.pos) >= despan_dist
      if away:
        self._world.remove(obj)

src/smartplay/crafter/crafter/env.py
microsoft/SmartPlay
import functools

import numpy as np
import opensimplex

from . import constants
from . import objects


def generate_world(world, player):
  simplex = opensimplex.OpenSimplex(seed=world.random.randint(0, 2 ** 31 - 1))
  tunnels = np.zeros(world.area, bool)
  for x in range(world.area[0]):
    for y in range(world.area[1]):
      _set_material(world, (x, y), player, tunnels, simplex)
  for x in range(world.area[0]):
    for y in range(world.area[1]):
      _set_object(world, (x, y), player, tunnels)


def _set_material(world, pos, player, tunnels, simplex):
  x, y = pos
  simplex = functools.partial(_simplex, simplex)
  uniform = world.random.uniform
  start = 4 - np.sqrt((x - player.pos[0]) ** 2 + (y - player.pos[1]) ** 2)
  start += 2 * simplex(x, y, 8, 3)
  start = 1 / (1 + np.exp(-start))
  water = simplex(x, y, 3, {15: 1, 5: 0.15}, False) + 0.1
  water -= 2 * start
  mountain = simplex(x, y, 0, {15: 1, 5: 0.3})
  mountain -= 4 * start + 0.3 * water
  if start > 0.5:
    world[x, y] = 'grass'
  elif mountain > 0.15:
    if (simplex(x, y, 6, 7) > 0.15 and mountain > 0.3):  # cave
      world[x, y] = 'path'
    elif simplex(2 * x, y / 5, 7, 3) > 0.4:  # horizonal tunnle
      world[x, y] = 'path'
      tunnels[x, y] = True
    elif simplex(x / 5, 2 * y, 7, 3) > 0.4:  # vertical tunnle
      world[x, y] = 'path'
      tunnels[x, y] = True
    elif simplex(x, y, 1, 8) > 0 and uniform() > 0.85:
      world[x, y] = 'coal'
    elif simplex(x, y, 2, 6) > 0.4 and uniform() > 0.75:
      world[x, y] = 'iron'
    elif mountain > 0.18 and uniform() > 0.994:
      world[x, y] = 'diamond'
    elif mountain > 0.3 and simplex(x, y, 6, 5) > 0.35:
      world[x, y] = 'lava'
    else:
      world[x, y] = 'stone'
  elif 0.25 < water <= 0.35 and simplex(x, y, 4, 9) > -0.2:
    world[x, y] = 'sand'
  elif 0.3 < water:
    world[x, y] = 'water'
  else:  # grassland
    if simplex(x, y, 5, 7) > 0 and uniform() > 0.8:
      world[x, y] = 'tree'
    else:
      world[x, y] = 'grass'


def _set_object(world, pos, player, tunnels):
  x, y = pos
  uniform = world.random.uniform
  dist = np.sqrt((x - player.pos[0]) ** 2 + (y - player.pos[1]) ** 2)
  material, _ = world[x, y]
  if material not in constants.walkable:
    pass
  elif dist > 3 and material == 'grass' and uniform() > 0.985:
    world.add(objects.Cow(world, (x, y)))
  elif dist > 10 and uniform() > 0.993:
    world.add(objects.Zombie(world, (x, y), player))
  elif material == 'path' and tunnels[x, y] and uniform() > 0.95:
    world.add(objects.Skeleton(world, (x, y), player))


def _simplex(simplex, x, y, z, sizes, normalize=True):
  if not isinstance(sizes, dict):
    sizes = {sizes: 1}
  value = 0
  for size, weight in sizes.items():
    if hasattr(simplex, 'noise3d'):
      noise = simplex.noise3d(x / size, y / size, z)
    else:
      noise = simplex.noise3(x / size, y / size, z)
    value += weight * noise
  if normalize:
    value /= sum(sizes.values())
  return value

src/smartplay/crafter/crafter/worldgen.py
microsoft/SmartPlay
import gym
from gym import error, spaces, utils
from gym.utils import seeding
from ..utils import HistoryTracker
from .crafter import Env
import numpy as np

id_to_item = [0]*19
import itertools
dummyenv = Env()
for name, ind in itertools.chain(dummyenv._world._mat_ids.items(), dummyenv._sem_view._obj_ids.items()):
    name = str(name)[str(name).find('objects.')+len('objects.'):-2].lower() if 'objects.' in str(name) else str(name)
    id_to_item[ind] = name
player_idx = id_to_item.index('player')
del dummyenv

vitals = ["health","food","drink","energy",]

rot = np.array([[0,-1],[1,0]])
directions = ['front', 'right', 'back', 'left']

def describe_inventory(info):
    result = ""
    
    status_str = "Your status:\n{}".format("\n".join(["- {}: {}/9".format(v, info['inventory'][v]) for v in vitals]))
    result += status_str + "\n\n"
    
    inventory_str = "\n".join(["- {}: {}".format(i, num) for i,num in info['inventory'].items() if i not in vitals and num!=0])
    inventory_str = "Your inventory:\n{}".format(inventory_str) if inventory_str else "You have nothing in your inventory."
    result += inventory_str #+ "\n\n"
    
    return result.strip()


REF = np.array([0, 1])

def rotation_matrix(v1, v2):
    dot = np.dot(v1,v2)
    cross = np.cross(v1,v2)
    rotation_matrix = np.array([[dot, -cross],[cross, dot]])
    return rotation_matrix

def describe_loc(ref, P):
    desc = []
    if ref[1] > P[1]:
        desc.append("north")
    elif ref[1] < P[1]:
        desc.append("south")
    if ref[0] > P[0]:
        desc.append("west")
    elif ref[0] < P[0]:
        desc.append("east")

    return "-".join(desc)


def describe_env(info):
    assert(info['semantic'][info['player_pos'][0],info['player_pos'][1]] == player_idx)
    semantic = info['semantic'][info['player_pos'][0]-info['view'][0]//2:info['player_pos'][0]+info['view'][0]//2+1, info['player_pos'][1]-info['view'][1]//2+1:info['player_pos'][1]+info['view'][1]//2]
    center = np.array([info['view'][0]//2,info['view'][1]//2-1])
    result = ""
    x = np.arange(semantic.shape[1])
    y = np.arange(semantic.shape[0])
    x1, y1 = np.meshgrid(x,y)
    loc = np.stack((y1, x1),axis=-1)
    dist = np.absolute(center-loc).sum(axis=-1)
    obj_info_list = []
    
    facing = info['player_facing']
    target = (center[0] + facing[0], center[1] + facing[1])
    target = id_to_item[semantic[target]]
    obs = "You face {} at your front.".format(target, describe_loc(np.array([0,0]),facing))
    
    for idx in np.unique(semantic):
        if idx==player_idx:
            continue

        smallest = np.unravel_index(np.argmin(np.where(semantic==idx, dist, np.inf)), semantic.shape)
        obj_info_list.append((id_to_item[idx], dist[smallest], describe_loc(np.array([0,0]), smallest-center)))

    if len(obj_info_list)>0:
        status_str = "You see:\n{}".format("\n".join(["- {} {} steps to your {}".format(name, dist, loc) for name, dist, loc in obj_info_list]))
    else:
        status_str = "You see nothing away from you."
    result += status_str + "\n\n"
    result += obs.strip()
    
    return result.strip()


def describe_act(info):
    result = ""
    
    action_str = info['action'].replace('do_', 'interact_')
    action_str = action_str.replace('move_up', 'move_north')
    action_str = action_str.replace('move_down', 'move_south')
    action_str = action_str.replace('move_left', 'move_west')
    action_str = action_str.replace('move_right', 'move_east')
    
    act = "You took action {}.".format(action_str) 
    result+= act
    
    return result.strip()


def describe_status(info):
    if info['sleeping']:
        return "You are sleeping, and will not be able take actions until energy is full.\n\n"
    elif info['dead']:
        return "You died.\n\n"
    else:
        return ""

    
def describe_frame(info, action):
    try:
        result = ""
        
        if action is not None:
            result+=describe_act(info)
        result+=describe_status(info)
        result+="\n\n"
        result+=describe_env(info)
        result+="\n\n"
        result+=describe_inventory(info)
        
        return result.strip()
    except:
        return "Error, you are out of the map."

class Crafter(Env):

    default_iter = 10
    default_steps = 10000

    def __init__(self, area=(64, 64), view=(9, 9), size=(64, 64), reward=True, length=10000, seed=None, max_steps=2):
        self.history = HistoryTracker(max_steps)
        self.action_list = ["Noop", "Move West", "Move East", "Move North", "Move South", "Do", \
    "Sleep", "Place Stone", "Place Table", "Place Furnace", "Place Plant", \
    "Make Wood Pickaxe", "Make Stone Pickaxe", "Make Iron Pickaxe", "Make Wood Sword", \
    "Make Stone Sword", "Make Iron Sword"]
        import pickle
        import pathlib
        root = pathlib.Path(__file__).parent
        with open(root / "assets/crafter_ctxt.pkl", 'rb') as f: # Context extracted using text-davinci-003 following https://arxiv.org/abs/2305.15486
            CTXT = pickle.load(f)
        CTXT = CTXT.replace("DO NOT answer in LaTeX.", "")
        CTXT = CTXT.replace("Move Up: Flat ground above the agent.", "Move North: Flat ground north of the agent.")
        CTXT = CTXT.replace("Move Down: Flat ground below the agent.", "Move South: Flat ground south of the agent.")
        CTXT = CTXT.replace("Move Left: Flat ground left to the agent.", "Move West: Flat ground west of the agent.")
        CTXT = CTXT.replace("Move Right: Flat ground right to the agent.", "Move East: Flat ground east of the agent.")
        self.desc = CTXT
        self.score_tracker = 0
        super().__init__(area, view, size, reward, length, seed)

    def reset(self):
        self.history.reset()
        super().reset()
        obs, reward, done, info = self.step(0)
        self.score_tracker = 0 + sum([1. for k,v in info['achievements'].items() if v>0])
        info.update({'manual': self.desc,
                'obs': describe_frame(info, None),
                'history': self.history.describe(),
                'score': self.score_tracker,
                'done': done,
                'completed': 0,
                })
        self.history.step(info)
        return obs, info
    
    def step(self, action):
        obs, reward, done, info = super().step(action)
        self.score_tracker = self.score_tracker + sum([1. for k,v in info['achievements'].items() if v>0])
        info.update({'manual': self.desc,
                'obs': describe_frame(info, self.action_list[action]),
                'history': self.history.describe(),
                'score': self.score_tracker,
                'done': done,
                'completed': 0,
                })
        self.history.step(info)
        return obs, reward, done, info

src/smartplay/crafter/crafter_env.py
microsoft/SmartPlay
# Benchmark games used in the paper
benchmark_games_v0 = [
    'RockPaperScissorBasic',
    'BanditTwoArmedHighLowFixed',
    'MessengerL1',
    'MessengerL2',
    'Hanoi3Disk',
    'Crafter',
    'MinedojoCreative0',
]


from . import recorded_settings, game_challenges
import pandas as pd
import numpy as np


def normalize_score(game_name, score):
    # Normalize score to [0, 1] for a given environment

    if game_name in recorded_settings.keys():
        return (score - recorded_settings[game_name]['min score']) / (recorded_settings[game_name]['human score'] - recorded_settings[game_name]['min score'])
    else:
        raise ValueError('Environment `{}` does not have recorded settings.'.format(game_name))


def analyze_capabilities(score_dict):
    # Analyze capabilities of a model based on the metric dictionary.
    # The metric dictionary should have the following format:
    # {
    #     'game_name': score,
    #     ...
    # }
    # The output is a dictionary with the following format:
    # {
    #     'game_name': capability,
    #     ...
    # }

    challenges_df = pd.DataFrame(game_challenges)
    challenges_df = challenges_df.loc[:, list(score_dict.keys())] - 1

    normalized_scores = np.array(list(score_dict.values())).reshape(-1, 1)
    scores = (challenges_df.values @ normalized_scores).squeeze() / np.sum(challenges_df.values, axis=1)

    return dict(zip(challenges_df.index, scores.flatten()))
src/smartplay/eval.py
microsoft/SmartPlay
from gym.envs.registration import register

from .hanoi_env import Hanoi3Disk, Hanoi4Disk

environments = [
    ['Hanoi3Disk', 'v0'],
    ['Hanoi4Disk', 'v0'],
]

for environment in environments:
    register(
        id='{}-{}'.format(environment[0], environment[1]),
        entry_point='smartplay.hanoi:{}'.format(environment[0]),
        # group='smartplay',
    )

src/smartplay/hanoi/__init__.py
microsoft/SmartPlay
import gym
from gym import error, spaces, utils
from gym.utils import seeding
from ..utils import HistoryTracker, describe_act

import random
import itertools
import numpy as np

action_to_move = [(0, 1), (0, 2), (1, 0),
                  (1, 2), (2, 0), (2, 1)]

class HanoiEnv(gym.Env):
    default_iter = 10
    default_steps = 30

    def __init__(self, max_steps=5, num_disks=4, env_noise=0):
        self.num_disks = num_disks
        self.env_noise = env_noise
        self.action_space = spaces.Discrete(6)
        self.observation_space = spaces.Tuple(self.num_disks*(spaces.Discrete(3),))

        self.current_state = None
        self.goal_state = self.num_disks*(2,)
        self.history = HistoryTracker(max_steps)

        self.done = None
        self.ACTION_LOOKUP = {0 : "(0,1) - top disk of pole 0 to top of pole 1 ",
                              1 : "(0,2) - top disk of pole 0 to top of pole 2 ",
                              2 : "(1,0) - top disk of pole 1 to top of pole 0",
                              3 : "(1,2) - top disk of pole 1 to top of pole 2",
                              4 : "(2,0) - top disk of pole 2 to top of pole 0",
                              5 : "(2,1) - top disk of pole 2 to top of pole 1"}


        self.action_list = [
        "Move top disk of rod A to top of rod B",
        "Move top disk of rod A to top of rod C",
        "Move top disk of rod B to top of rod A",
        "Move top disk of rod B to top of rod C",
        "Move top disk of rod C to top of rod A",
        "Move top disk of rod C to top of rod B",
        ]
      
        self.desc = """
The game consists of three rods (A,B,C) and a number of disks of various sizes, which can go onto any rod. 
The game begins with the disks stacked on rod A in order of decreasing size, the smallest at the top (righthand side). 
The objective is to move the entire stack to rod C, obeying the following rules:

 - Only one disk may be moved at a time.
 - Each move consists of taking the top disk from one of the stacks and placing it on top of another stack or on an empty rod.
 - You cannot place a bigger disk on top of a smaller disk.

For example, considering movements from B under the following setting:
- A: |bottom, [0], top|
- B: |bottom, [1], top|
- C: |bottom, [2], top|
You are only allowed to move from B to C but not A, since the top of B (1) is smaller than the top of C (2) but bigger than the top of A (0).

Finally, the starting configuration is:
- A: |bottom, {}, top|
- B: |bottom, [], top|
- C: |bottom, [], top|

and the goal configuration is:
- A: |bottom, [], top|
- B: |bottom, [], top|
- C: |bottom, {}, top|
with top on the right and bottom on the left

{}
""".format(list(range(num_disks))[::-1],list(range(num_disks))[::-1], describe_act(self.action_list)).strip()

    def step(self, action):
        """
        * Inputs:
            - action: integer from 0 to 5 (see ACTION_LOOKUP)
        * Outputs:
            - current_state: state after transition
            - reward: reward from transition
            - done: episode state
            - info: dict of booleans (noisy?/invalid action?)
        0. Check if transition is noisy or not
        1. Transform action (0 to 5 integer) to tuple move - see Lookup
        2. Check if move is allowed
        3. If it is change corresponding entry | If not return same state
        4. Check if episode completed and return
        """
        if self.done:
            raise RuntimeError("Episode has finished. Call env.reset() to start a new episode.")

        info = {"transition_failure": False,
                "invalid_action": False}

        if self.env_noise > 0:
            r_num = random.random()
            if r_num <= self.env_noise:
                action = random.randint(0, self.action_space.n-1)
                info["transition_failure"] = True

        move = action_to_move[action]

        if self.move_allowed(move):
            disk_to_move = min(self.disks_on_peg(move[0]))
            moved_state = list(self.current_state)
            moved_state[disk_to_move] = move[1]
            self.current_state = tuple(moved_state)
        else:
            info["invalid_action"] = True

        if self.current_state == self.goal_state:
            reward = 100
            self.done = True
        elif info["invalid_action"] == True:
            reward = -1
        else:
            reward = 0

        info["state"] = (self.disks_on_peg(0), self.disks_on_peg(1), self.disks_on_peg(2))
        info["score"] = len(self.disks_on_peg(2))
        info["manual"] = self.desc
        info["obs"] = self.describe_state(info, action)
        info["history"] = self.history.describe()
        info["completed"] = 0 if not self.done else 1
        self.history.step(info)

        return self.current_state, reward, self.done, info

    def describe_state(self, state, action=None):
        index = ['A', 'B', 'C']
        if action is not None:
            result = "You tried to {}. Current configuration:".format(self.action_list[action].lower())
        else:
            result = "Current configuration:"
        for i in range(3):
            result += "\n- {}: |bottom, {}, top|".format(index[i], state['state'][i][::-1])
        return result.strip()

    def disks_on_peg(self, peg):
        """
        * Inputs:
            - peg: pole to check how many/which disks are in it
        * Outputs:
            - list of disk numbers that are allocated on pole
        """
        return [disk for disk in range(self.num_disks) if self.current_state[disk] == peg]

    def move_allowed(self, move):
        """
        * Inputs:
            - move: tuple of state transition (see ACTION_LOOKUP)
        * Outputs:
            - boolean indicating whether action is allowed from state!
        move[0] - peg from which we want to move disc
        move[1] - peg we want to move disc to
        Allowed if:
            * discs_to is empty (no disc of peg) set to true
            * Smallest disc on target pole larger than smallest on prev
        """
        disks_from = self.disks_on_peg(move[0])
        disks_to = self.disks_on_peg(move[1])

        if disks_from:
            return (min(disks_to) > min(disks_from)) if disks_to else True
        else:
            return False

    def reset(self):
        self.current_state = self.num_disks * (0,)
        self.done = False
        self.history.reset()
        info = {"state":(self.disks_on_peg(0), self.disks_on_peg(1), self.disks_on_peg(2))}
        info["score"] = len(self.disks_on_peg(2))
        info["manual"] = self.desc
        info["obs"] = self.describe_state(info)
        info["history"] = self.history.describe()
        info["completed"] = 0
        self.history.step(info)
        return self.current_state, info

    def render(self, mode='human', close=False):
        return

    def set_env_parameters(self, num_disks=4, env_noise=0, verbose=True):
        self.num_disks = num_disks
        self.env_noise = env_noise
        self.observation_space = spaces.Tuple(self.num_disks*(spaces.Discrete(3),))
        self.goal_state = self.num_disks*(2,)

        if verbose:
            print("Hanoi Environment Parameters have been set to:")
            print("\t Number of Disks: {}".format(self.num_disks))
            print("\t Transition Failure Probability: {}".format(self.env_noise))

    def get_movability_map(self, fill=False):
        # Initialize movability map
        mov_map = np.zeros(self.num_disks*(3, ) + (6,))

        if fill:
            # Get list of all states as tuples
            id_list = self.num_disks*[0] + self.num_disks*[1] + self.num_disks*[2]
            states = list(itertools.permutations(id_list, self.num_disks))

            for state in states:
                for action in range(6):
                    move = action_to_move[action]
                    disks_from = []
                    disks_to = []
                    for d in range(self.num_disks):
                        if state[d] == move[0]: disks_from.append(d)
                        elif state[d] == move[1]: disks_to.append(d)

                    if disks_from: valid = (min(disks_to) > min(disks_from)) if disks_to else True
                    else: valid = False

                    if not valid: mov_map[state][action] = -np.inf

                    move_from = [m[0] for m in action_to_move]
                    move_to = [m[1] for m in action_to_move]

        return mov_map

class Hanoi3Disk(HanoiEnv):
    """Basic 3 disk Hanoi Environment"""
    def __init__(self):
        HanoiEnv.__init__(self, num_disks=3)

class Hanoi4Disk(HanoiEnv):
    """Basic 4 disk Hanoi Environment"""
    def __init__(self):
        HanoiEnv.__init__(self, num_disks=4)

src/smartplay/hanoi/hanoi_env.py
microsoft/SmartPlay
from gym.envs.registration import register

from .messenger_env import MessengerEnv

environments = [
    ['MessengerL1', 'v0'],
    ['MessengerL2', 'v0'],
    ['MessengerL3', 'v0'],
]

for environment in environments:
    register(
        id='{}-{}'.format(environment[0], environment[1]),
        entry_point='smartplay.messenger_emma:MessengerEnv',
        kwargs={'lvl': int(environment[0][-1])}
        # group='smartplay',
    )

src/smartplay/messenger_emma/__init__.py
microsoft/SmartPlay
import gym
from gym import error, spaces, utils
from gym.utils import seeding
from ..utils import HistoryTracker, describe_act

import random
import itertools
import numpy as np
import messenger
import messenger.envs


''' Format function passed to numpy print to make things pretty.
'''

id_map = {}

for ent in messenger.envs.config.ALL_ENTITIES:
    id_map[ent.id] = ent.name
id_map[0] = '  '
id_map[15] = 'you (agent) without the message'
id_map[16] = 'you (agent) with the message'

def describe_block(i):
    if i < 17:
        return id_map[i]
    else:
        return 'XX'

def describe_loc(ref, P):
    desc = []
    if ref[0] > P[0]:
        desc.append("north")
    elif ref[0] < P[0]:
        desc.append("south")
    if ref[1] > P[1]:
        desc.append("west")
    elif ref[1] < P[1]:
        desc.append("east")

    return "-".join(desc)

    
def describe_frame(info):

    if 15 in np.unique(info['avatar']):
        obs = "You (agent) don't have the message."
        agent = 15
    elif 16 in np.unique(info['avatar']):
        obs = "You (agent) already have the message."
        agent = 16
    else:
        print(np.unique(info['avatar']))
        raise NotImplemented("Problem with agent")
    center = np.array(np.where(info['avatar'].squeeze() == agent)).squeeze()
    info = info['entities']
    result = ""
    x = np.arange(info.shape[1])
    y = np.arange(info.shape[0])
    x1, y1 = np.meshgrid(x,y)
    loc = np.stack((y1, x1),axis=-1)
    dist = np.absolute(center-loc).sum(axis=-1)[:,:,np.newaxis]

    obj_info_list = []
    
    
    for idx in np.unique(info):
        if idx == 15 or idx == 0:
            continue
        smallest = np.unravel_index(np.argmin(np.where(info==idx, dist, np.inf)), info.shape)
        obj_info_list.append((describe_block(idx), dist[smallest[:2]][0], describe_loc(np.array([0,0]), smallest[:2]-center)))

    if len(obj_info_list)>0:
        status_str = "You see:\n{}".format("\n".join(["- {} {} steps to your {}".format(name, dist, loc) if dist>0 else "- {} 0 steps with you".format(name) for name, dist, loc in obj_info_list]))
    else:
        status_str = "You see nothing away from you."
    result += obs.strip() + "\n\n" + status_str.strip()
    
    return result.strip()

class MessengerEnv(gym.Env):
    default_iter = 100
    default_steps = None

    def __init__(self, lvl=1, max_steps=2, env_noise=0):

        lvl_to_steps = [4, 64, 128]
        self.default_steps = lvl_to_steps[lvl-1]
        env_id = 'msgr-test-v{}'.format(lvl)
        self._env = gym.make(env_id)
        self.action_list = ["Move North", "Move South", "Move West", "Move East", "Do Nothing"]

        self.history = HistoryTracker(max_steps)
        self.game_context = """In the game, MESSENGER, each entity can take on one of three roles: an enemy, message, or goal. The agent’s objective is to bring the message to the goal while avoiding the enemies. If the agent encounters an enemy at any point in the game, or the goal without first obtaining the message, it loses the game and obtains a reward of −1.""".strip()
        self.advice = """To solve a game, you may find it helpful to list the objects that you see. Then for each object, match it with an entity description, and identify whether it is good or bad to interact with the object.
The name specifications of in-game objects may not be exact matches. Please try identifying with synonyms.
""".strip()

    def _update_manual(self, manual):
        self.desc = "{}\n\n{}\n\n{}\n\n{}".format(self.game_context, "\n".join(manual), self.advice, describe_act(self.action_list)).strip()

    def describe(self, obs, action=None):
        if action is not None:
            return "You took action {}.\n\n{}".format(self.action_list[action], describe_frame(obs))
        else:
            return describe_frame(obs)

    def reset(self):
        obs, manual = self._env.reset()
        self._update_manual(manual)
        self.history.reset()
        info = {
            "obs": self.describe(obs),
            "manual": self.desc,
            "history": self.history.describe(),
            "score": 0,
            "completed": 0,
            }
        self.history.step(info)
        return obs, info

    def step(self, action):
        obs, reward, done, info = self._env.step(action)
        try:
            description = self.describe(obs, action)
        except:
            description = "Environment Error."
        info={
            "obs": description,
            "manual": self.desc,
            "history": self.history.describe(),
            "score": reward,
            "completed": int(reward==1 and done),
            }
        self.history.step(info)
        return obs, reward, done, info

src/smartplay/messenger_emma/messenger_env.py
microsoft/SmartPlay
from gym.envs.registration import register

from .minedojo_env import MineDojoEnv

environments = [
    ['MinedojoCreative0', 'v0'],
    ['MinedojoCreative1', 'v0'],
    ['MinedojoCreative2', 'v0'],
    ['MinedojoCreative4', 'v0'],
    ['MinedojoCreative5', 'v0'],
    ['MinedojoCreative7', 'v0'],
    ['MinedojoCreative8', 'v0'],
    ['MinedojoCreative9', 'v0'],
]

for environment in environments:
    register(
        id='{}-{}'.format(environment[0], environment[1]),
        entry_point='smartplay.minedojo:MineDojoEnv',
        kwargs={'task_id': environment[0][-1]}
    )

src/smartplay/minedojo/__init__.py
microsoft/SmartPlay
import gym
from gym import error, spaces, utils
from gym.utils import seeding
from ..utils import HistoryTracker, describe_act

import random
import itertools
import numpy as np
import math
import cc3d
import minedojo

biomes_dict = {
    0: 'Ocean',
    1: 'Plains',
    2: 'Desert',
    3: 'Extreme Hills',
    4: 'Forest',
    5: 'Taiga',
    6: 'Swampland',
    7: 'River',
    8: 'Hell (The Nether)',
    9: 'The End',
    10: 'FrozenOcean',
    11: 'FrozenRiver',
    12: 'Ice Plains',
    13: 'Ice Mountains',
    14: 'MushroomIsland',
    15: 'MushroomIslandShore',
    16: 'Beach',
    17: 'DesertHills',
    18: 'ForestHills',
    19: 'TaigaHills',
    20: 'Extreme Hills Edge',
    21: 'Jungle',
    22: 'JungleHills',
    23: 'JungleEdge',
    24: 'Deep Ocean',
    25: 'Stone Beach',
    26: 'Cold Beach',
    27: 'Birch Forest',
    28: 'Birch Forest Hills',
    29: 'Roofed Forest',
    30: 'Cold Taiga',
    31: 'Cold Taiga Hills',
    32: 'Mega Taiga',
    33: 'Mega Taiga Hills',
    34: 'Extreme Hills+',
    35: 'Savanna',
    36: 'Savanna Plateau',
    37: 'Mesa',
    38: 'Mesa Plateau F',
    39: 'Mesa Plateau',
    127: 'The Void',
    129: 'Sunflower Plains',
    130: 'Desert M',
    131: 'Extreme Hills M',
    132: 'Flower Forest',
    133: 'Taiga M',
    134: 'Swampland M',
    140: 'Ice Plains Spikes',
    149: 'Jungle M',
    151: 'JungleEdge M',
    155: 'Birch Forest M',
    156: 'Birch Forest Hills M',
    157: 'Roofed Forest M',
    158: 'Cold Taiga M',
    160: 'Mega Spruce Taiga',
    161: 'Redwood Taiga Hills M',
    162: 'Extreme Hills+ M',
    163: 'Savanna M',
    164: 'Savanna Plateau M',
    165: 'Mesa (Bryce)',
    166: 'Mesa Plateau F M',
    167: 'Mesa Plateau M'
}

yaw_granularity = 6
pitch_granularity = 6
FOV = 96

pitch_cnt = len(np.arange(-FOV//2, FOV//2+1, pitch_granularity))
yaw_cnt = len(np.arange(-FOV//2, FOV//2+1, yaw_granularity))

def get_direction(yaw, pitch):
    # Convert yaw and pitch to radians
    yaw_rad = math.radians(yaw)
    pitch_rad = math.radians(pitch)
    direction_list = []

    # Calculate the x, y, and z components of the direction vector
    x = math.sin(yaw_rad)* math.cos(pitch_rad)
    y = math.sin(pitch_rad)
    z = math.cos(yaw_rad) * math.cos(pitch_rad)

    # Determine the direction based on the sign of the x, y, and z components
    if y > 0.2:
        direction = "above you"
    elif y < -0.2:
        direction = "below you"
    else:
        direction = "at your level"

    if z > 0.2:
        direction_list.append("north")
    elif z < -0.2:
        direction_list.append("south")
    
    if x > 0.2:
        direction_list.append("east")
    elif x < -0.2:
        direction_list.append("west")

    return direction + " to " + ("front" if "-".join(direction_list) == "" else "-".join(direction_list))

def describe_surround(obs):
    # print("What blocks are around me?")
    # print()
    result = "Around you:\n"
    
    # Voxel size (Can be changed based on the voxel space)
    voxel_x_size = 9
    voxel_y_size = 9
    voxel_z_size = 9
    voxel_center_x = voxel_x_size // 2
    voxel_center_y = voxel_y_size // 2
    voxel_center_z = voxel_z_size // 2
    
    # Convert the 3D block name list into np array
    block_names = np.array(obs['voxels']['block_name'])
    
    # Only preserve blocks 1 meter below me
    block_names = block_names[:, voxel_z_size // 2 - 1:, :]

    # Encode the block name using its unique index
    unique_block_names, unique_block_indices = np.unique(block_names, return_inverse=True)
    block_names_int = np.array(unique_block_indices).reshape(
        block_names.shape)

    # Apply 3D connected component and extract the number of labels
    block_names_labels, labels_count = cc3d.connected_components(block_names_int, return_N = True, connectivity = 6)

    # Get the centroids of each label
    centroids = cc3d.statistics(block_names_labels)['centroids']

    # Get the map: label -> block name
    label_2_block_name = [""] * (labels_count + 1)
    for i in range(voxel_x_size):
        for j in range(block_names.shape[1]):
            for k in range(voxel_y_size):
                label = block_names_labels[i][j][k]
                block_name = block_names[i][j][k]
                label_2_block_name[label] = block_name

    # Describe the surrounding environment based on the connected component labels
    # Helper lists for printing
    x_axis_list = ["west", "east"]
    y_axis_list = ["north", "south"]
    z_axis_list = ["below you", "above you"] # The height of me is defined as the height of my feet
    
    # The map: block name -> min distance
    block2dist = dict()
    
    # The map: block name -> direction
    block2dir = dict()
    
    # Get the direction of each label based on the centroid
    for i in range(labels_count + 1):
        middle_point = centroids[i]
        block_name = label_2_block_name[i]

        # Skip the air
        if block_name == 'air': continue

        # Get the indices to the printing help list
        x_coord = 1 if middle_point[0] > voxel_center_x else 0
        z_coord = 1 if middle_point[1] > 1 else 0
        y_coord = 1 if middle_point[2] > voxel_center_y else 0

        is_x_center = math.isclose(middle_point[0], voxel_center_x)
        is_z_center = math.isclose(middle_point[1], 1)
        is_y_center = math.isclose(middle_point[2], voxel_center_y)

        # If the component is exactly centered, skip
        if is_x_center and is_y_center and is_z_center:
            continue
        
        # Update the min distance and correspoding direction for the same block name
        distance = math.sqrt((middle_point[0] - voxel_center_x)**2 + (middle_point[1] - 1)**2 + (middle_point[2] - voxel_center_y)**2)
        if block_name not in block2dist or distance < block2dist[block_name]:
            block2dist[block_name] = distance
            
            # Record the direction of this component
            direction_str_list = []

            if not is_z_center:
                direction_str_list.append(z_axis_list[z_coord])
            else:
                direction_str_list.append("")

            if not is_y_center:
                direction_str_list.append(y_axis_list[y_coord])

            # Skip when the coordinate is at the voxel center
            if not is_x_center:
                direction_str_list.append(x_axis_list[x_coord])

            direction_str = direction_str_list[0] + " to "+ ('-'.join(direction_str_list[1:]) if '-'.join(direction_str_list[1:]) != "" else "front")
            
            block2dir[block_name] = direction_str.strip()
        
    # Finally print the blocks
    for block_name in block2dist:
        direction_str = block2dir[block_name]
        result += f" - {block_name}, {'%.2f' % block2dist[block_name]} blocks away, {direction_str}\n"
        
    return result.strip()

# Describe the block the agent is facing to

def describe_cursor(obs):
    
    # Track the cursor using lidar
    cursor_block = obs["rays"]["block_name"][0]
    
    # Skip if the cursor is pointing to air
    if cursor_block == 'air':
        return "You're not aiming at any block."
    else:
        return f"You're aiming at {cursor_block}."

# Describe the surrounding entities around the agent

def describe_entity(obs):
    result = ""
    my_yaw = obs["location_stats"]["yaw"]
    my_pitch = obs["location_stats"]["pitch"]
    
    # Number of pitch/yaw rays (Can be changed based on the lidar rays)
    
    # A flag indicating where there are entities around
    see_entity = False

    # Reshape and convert the list into np array
    entity_names = np.array(obs["rays"]["entity_name"][1:].reshape(pitch_cnt,yaw_cnt))

    # Encode the block name using its unique index
    unique_entity_names, unique_entity_indices = np.unique(entity_names, return_inverse=True)
    entity_names_int = np.array(unique_entity_indices).reshape(entity_names.shape)

    # Apply 2D connected components
    entity_names_labels, labels_count = cc3d.connected_components(entity_names_int, return_N = True, connectivity = 8)

    # Get the map: label -> entity name
    label_2_entity_name = [""] * (labels_count + 1)
    for i in range(pitch_cnt):
        for j in range(yaw_cnt):
            label = entity_names_labels[i][j]
            entity_name = entity_names[i][j]
            label_2_entity_name[label] = entity_name

    # Describe each component
    for i in range(labels_count + 1):
        entity_name = label_2_entity_name[i]

        # Skip the null entity
        if entity_name == "null": 
            continue
        else: 
            see_entity = True

        # Find all the indices of this component
        all_idx = np.where(entity_names_labels == i)
        amount = np.sum(entity_names_labels == i)

        # Find the minimum distance inside each component
        min_distance = math.inf
        direction = ""
        for row, col in zip(all_idx[0], all_idx[1]):
            index = row * yaw_cnt + col
            distance = obs["rays"]["entity_distance"][index+1]
            if distance < min_distance:
                min_distance = distance
                yaw = ((col / yaw_cnt) * FOV - FOV/2 + my_yaw + 540) % 360 - 180
                pitch = (row / pitch_cnt) * 180 - 90 + my_pitch
                direction = get_direction(yaw, pitch)

        amount_description = "taking {0:.0f}% of screen".format(amount/(pitch_cnt*yaw_cnt)*100)

        result+=f" - {entity_name}, {'%.2f' % min_distance} blocks away, {direction}, {amount_description}\n"
        
    # If there are no surrounding entities that the agent can see
    if not see_entity:
        result = ""
    
    return result

# Describe the surrounding entities around the agent

def describe_obj(obs):
    result = "You see:\n"
    my_yaw = obs["location_stats"]["yaw"]
    my_pitch = obs["location_stats"]["pitch"]
    
    # Number of pitch/yaw rays (Can be changed based on the lidar rays)
    
    # A flag indicating where there are entities around
    see_entity = False

    # Reshape and convert the list into np array
    entity_names = np.array(obs["rays"]["block_name"][1:].reshape(pitch_cnt,yaw_cnt))

    # Encode the block name using its unique index
    unique_entity_names, unique_entity_indices = np.unique(entity_names, return_inverse=True)
    entity_names_int = np.array(unique_entity_indices).reshape(entity_names.shape)

    # Apply 2D connected components
    entity_names_labels, labels_count = cc3d.connected_components(entity_names_int, return_N = True, connectivity = 8)

    # Get the map: label -> entity name
    label_2_entity_name = [""] * (labels_count + 1)
    for i in range(pitch_cnt):
        for j in range(yaw_cnt):
            label = entity_names_labels[i][j]
            entity_name = entity_names[i][j]
            label_2_entity_name[label] = entity_name

    # Describe each component
    for i in range(labels_count + 1):
        entity_name = label_2_entity_name[i]

        # Skip the null entity
        if entity_name in ("null", "air"): 
            continue
        else: 
            see_entity = True

        # Find all the indices of this component
        all_idx = np.where(entity_names_labels == i)
        amount = np.sum(entity_names_labels == i)

        # Find the minimum distance inside each component
        min_distance = math.inf
        direction = ""
        for row, col in zip(all_idx[0], all_idx[1]):
            index = row * yaw_cnt + col
            distance = obs["rays"]["block_distance"][index+1]
            if distance < min_distance:
                min_distance = distance
                yaw = ((col / yaw_cnt) * FOV - FOV/2 + my_yaw + 540) % 360 - 180
                pitch = ((row / pitch_cnt) * 180 - 90) + my_pitch
                direction = get_direction(yaw, pitch)
        
        amount_description = "taking {0:.0f}% of screen".format(amount/(pitch_cnt*yaw_cnt)*100)

        result+=f" - {entity_name}, {'%.2f' % min_distance} blocks away, {direction}, {amount_description}\n"
        
    # If there are no surrounding entities that the agent can see
    if not see_entity:
        result = ""

    return result.strip()

# Describe the exact coordinate the agent is in
def describe_location(obs):
    result = ""

    coord_list = obs["location_stats"]["pos"]
    
    # Describe the direction the agent is currently facing
    yaw = obs["location_stats"]["yaw"]
    
    direction_list = ["north", "north-east", "east", "south-east", "south", "south-west", "west", "north-west"]
    direction_index = (int((yaw // 22.5) + 1)) // 2 # Calculate the index mapping to the direction_list above
    if direction_index == 4:
        direction_index *= -1
    direction_index += 4

    result = f"Coordinate ({'%.2f' % coord_list[0]},{'%.2f' % coord_list[1]},{'%.2f' % coord_list[2]}). Facing {direction_list[direction_index]}."    # Describe whether the agent is looking up or down
    pitch = obs["location_stats"]["pitch"]

    if pitch < 0:
        result += " Looking up."
    elif pitch > 0:
        result += " Looking down."

    return result.strip()

def describe_frame(obs):
    return describe_location(obs)+"\n\n"+describe_cursor(obs)+"\n"+describe_surround(obs)+"\n"+describe_entity(obs)+"\n"+describe_obj(obs)

class MineDojoEnv(gym.Env):
    default_iter = 20
    default_steps = 200

    def __init__(self, task_id, max_steps=2):

        # Creative task
        self._env = minedojo.make(
            task_id="creative:{}".format(task_id),
            fast_reset=False,
            image_size=(160, 256),
            use_voxel=True,
            voxel_size=dict(xmin=-4, ymin=-4, zmin=-4, xmax=4, ymax=4, zmax=4), # Set voxel space to be 9*9*9
            use_lidar=True,
            lidar_rays=[(0,0,3)] + [
                    (np.pi * pitch / 180, np.pi * yaw / 180, 65)
                    for pitch in np.arange(-FOV//2, FOV//2+1, pitch_granularity)
                    for yaw in np.arange(-FOV//2, FOV//2+1, yaw_granularity)
            ]   # Track the agent's cursor
        )
        goal = self._env.task_prompt.replace(".", "")[self._env.task_prompt.find('find ')+len('find '):]
        self.goal_set = {k for k,v in biomes_dict.items() if goal in v.lower()}
        self.action_list = ["Move North", "Move East", "Move South", "Move West"]
        self.noop_step=50

        self.history = HistoryTracker(max_steps)
        self.desc = """
You are in Minecraft and your goal is to find a {} biome. You are not allowed to craft anything.

In your observation, you are provided the amount of space an object takes in your field of view. Note that objects of the same size takes more space when they are closer to you.

{}
""".format(goal, describe_act(self.action_list)).strip()
        
    def get_turn_action(self, act, obs):
        yaw = obs["location_stats"]["yaw"]
        # Compute yaw delta to face the yaw indicated by the action 'act'
        action_yaw = act * 90 - 180
        yaw_delta = (action_yaw - yaw + 180) % 360 // 15
        action = self._env.action_space.no_op()
        action[4] = yaw_delta
        return action

    def go_forward_action(self):
        action = self._env.action_space.no_op()
        action[0] = 1 # Move forward
        action[1] = 0 if random.random() < 0.5 else (1 if random.random() < 0.5 else 2) # Move left/right with prob=0.5 to get unstuck
        action[2] = 3 if random.random() < 0.5 else 1 # Sprint with prob=0.5
        return action

    def describe(self, obs, action=None):
        if action is not None:
            return "You took action {}.\n\n{}".format(self.action_list[action], describe_frame(obs))
        else:
            return describe_frame(obs)

    def reset(self):
        obs = self._env.reset()
        self.history.reset()
        info = {
            "obs": self.describe(obs),
            "manual": self.desc,
            "history": self.history.describe(),
            "score": 0,
            "completed": 0,
            }
        self.history.step(info)
        self.obs = obs
        return obs, info

    def step(self, action):
        R = 0
        obs, reward, done, info = self._env.step(self.get_turn_action(action, self.obs))
        R += reward
        for _ in range(self.noop_step):
            obs, reward, done, info = self._env.step(self.go_forward_action())
            R += reward
            if obs["location_stats"]["biome_id"].item() in self.goal_set:
                break

        info.update({
            "obs": self.describe(obs, action),
            "manual": self.desc,
            "history": self.history.describe(),
            "score": 0,
            "completed": 0,
            })
        if obs["location_stats"]["biome_id"].item() in self.goal_set:
            info["completed"] = 1
            info["score"] = 1
            done = True
        self.history.step(info)
        self.obs = obs
        return obs, R, done, info

src/smartplay/minedojo/minedojo_env.py
microsoft/SmartPlay
from gym.envs.registration import register

from .rock_paper_scissor import RockPaperScissorBasic
from .rock_paper_scissor import RockPaperScissorDifferentScore

environments = [
    ['RockPaperScissorBasic', 'v0'],
    ['RockPaperScissorDifferentScore', 'v0'],
]

for environment in environments:
    register(
        id='{}-{}'.format(environment[0], environment[1]),
        entry_point='smartplay.rock_paper_scissors:{}'.format(environment[0]),
        # group='smartplay',
    )

src/smartplay/rock_paper_scissors/__init__.py
microsoft/SmartPlay
import numpy as np
import gym
from gym import spaces
from gym.utils import seeding
from ..utils import HistoryTracker, describe_act
import random


class RPSEnv(gym.Env):
    """
    Bandit environment base to allow agents to interact with the class n-armed bandit
    in different variations

    p_dist:
        A list of probabilities of the likelihood that a particular bandit will pay out
    r_dist:
        A list of either rewards (if number) or means and standard deviations (if list)
        of the payout that bandit has
    """

    default_iter = 20
    default_steps = 50

    def __init__(self, probs, reward = [1, 1, 1], max_steps=50):

        if min(probs) < 0 or max(probs) > 1 or sum(probs) != 1:
            raise ValueError("All probabilities must be between 0 and 1, and sum to 1")

        dist = list(zip(probs, reward))
        random.shuffle(dist)
        self.probs = [d[0] for d in dist]

        self.action_list = ["Rock", "Paper", "Scissor"]
        self.reward = [d[1] for d in dist]
        self._update_manual()
        self.observation_space = spaces.Discrete(1)
        self.action_space = spaces.Discrete(len(self.probs))
        self.history = HistoryTracker(max_steps)
        self.score_tracker = 0

        self._seed()

    def _update_manual(self):
        self.desc = """
For the game Rock Paper Scissors, you and the opponent choose one of three options: rock, paper, or scissors.
After both players have chosen, the winner is determined as follows:
Rock crushes scissors (Rock wins, score {})
Scissors cut paper (Scissors win, score {})
Paper covers rock (Paper wins, score {})
If you lose, your score is the negative of the winner's score.
If both players choose the same option, it's a draw (score 0).
Your goal is to maximize your score.

{}
""".strip().format(*self.reward, describe_act(self.action_list))

    def compute_optimal_action(self):
        n_actions = len(self.probs)
        expected_scores = np.zeros(n_actions)

        for action in range(n_actions):
            expected_score = 0
            for opponent_action in range(n_actions):
                expected_score += self.probs[opponent_action] * self.reward[opponent_action]
                if (action == 0 and opponent_action == 1) or (action == 1 and opponent_action == 2) or (action == 2 and opponent_action == 0):
                    expected_score -= self.probs[opponent_action] * self.reward[opponent_action]
                elif action == opponent_action:
                    pass
                else:
                    expected_score += self.probs[opponent_action] * self.reward[action]
            expected_scores[action] = expected_score

        optimal_action = np.argmax(expected_scores)
        expected_score = expected_scores[optimal_action]

        return optimal_action, expected_score

    def sample_opponent_action(self, probs):
        return np.random.choice(len(probs), p=probs)

    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, action):
        assert self.action_space.contains(action)
        opponent_action = self.sample_opponent_action(self.probs)
        
        if action == 0 and opponent_action == 1:
            result = "lost"
        elif action == 1 and opponent_action == 2:
            result = "lost"
        elif action == 2 and opponent_action == 0:
            result = "lost"
        elif action == opponent_action:
            result = "tied"
        else:
            result = "won"

        reward = self.reward[action] if result == "won" else -self.reward[opponent_action] if result == "lost" else 0
        done = False
        optimal_action, expected_score = self.compute_optimal_action()
        self.score_tracker += int(action == optimal_action)
        info = {
            "obs": "You chose {}, and the opponent chose {}. You {} and received score {}.\nNew round begins.".format(
                self.action_list[action], 
                self.action_list[opponent_action], 
                result, reward),
            "manual": self.desc,
            "history": self.history.describe(),
            "score": self.score_tracker,
            "result": result,
            "completed": 0,
            }
        self.history.step(info)

        return 0, reward, done, info

    def reset(self):
        probs, reward = self.probs, self.reward
        dist = list(zip(probs, reward))
        random.shuffle(dist)
        self.probs = [d[0] for d in dist]

        self.reward = [d[1] for d in dist]
        self._update_manual()
        self.history.reset()
        self.score_tracker = 0
        info = {
            "obs":"New round begins.",
            "manual": self.desc,
            "history": self.history.describe(),
            "score": self.score_tracker,
            "completed": 0,
            }
        self.history.step(info)
        return 0, info

    def render(self, mode='human', close=False):
        pass


class RockPaperScissorBasic(RPSEnv):
    """Basic version of the game with biased opponent"""
    def __init__(self):
        RPSEnv.__init__(self, [0.2, 0.2, 0.6])


class RockPaperScissorDifferentScore(RPSEnv):
    """Biased opponent and different scores"""
    def __init__(self):
        RPSEnv.__init__(self, p_dist=[0.3, 0.3, 0.4], r_dist=[1, 2, 1])

src/smartplay/rock_paper_scissors/rock_paper_scissor.py
microsoft/SmartPlay
import copy
class HistoryTracker:

    def __init__(self, max_steps) -> None:
        self.max_steps = max_steps
        self.game_step = 0
        self.reset()

    def step(self, info) -> None:
        self.info.append(copy.copy(info))
        if len(self.info) > self.max_steps:
            self.info.pop(0)
        self.game_step += 1

    def reset(self) -> None:
        self.info = []
        self.game_step = 0
    
    def describe(self, game_step=None):
        if len(self.info) == 0:
            return ""
        game_step = self.game_step if game_step is None else game_step
        result = "Most recent {} steps of the player's in-game observation:\n\n".format(len(self.info))
        for i, info in enumerate(self.info):
            result += "Player Observation Step {}:\n".format(game_step - len(self.info) + i)
            result += info["obs"] + "\n\n"
        return result.strip()
    
    def score(self):
        return sum([info["score"] for info in self.info])


def describe_act(action_list):
    return "List of all actions:\n" + "\n".join(["{}. {}".format(i+1, s) for i,s in enumerate(action_list)])

src/smartplay/utils.py
microsoft/SmartPlay