Nested spaces in gRPC

[sogartar]

❓ Questions and Help

It looks like that the proto definitions of action space, action, observation space and observation do not support nesting in.
On the other hand OpenAI Gym's interface has Dict and by extension CompilerGym has also Dict. That makes deep hierarchies possible there.

Does It make sense to extend the gRPC interface to support this type of structure?

Additional Context

This is not a very pressing issue as currently the structure I want to represent is simple enough, that I can use flat names like a.b
and a.c.
In the future there may be the need for more complex structures.

[ChrisCummins]

Hi @sogartar,

Does It make sense to extend the gRPC interface to support this type of structure?

Interesting suggestion! Possibly 🙂 I can see the value of nested action and observation spaces, but I don't know of any immediate use cases so I haven't prioritized it. The closest we have in the current version would be something like the JSON data structure returned by some of the LLVM environment's observations, but in that case I just send it over as raw bytes data and deserialize on the frontend.

Cheers,
Chris

[sogartar]

Currently the structure for the action spaces is

message ActionSpace {
  string name = 1;
  repeated ChoiceSpace choice = 2;
  bool named_choices = 3;
}

message ChoiceSpace {
  string name = 1;
  oneof space {
    ...
  }
}

I could extend the protocol by adding to the ChoiceSpace. Another option is to restructure it. I was thinking to make it closer to what is available in python.
Something like shown below.

package compiler_gym.spaces;

message ScalarLimit {
  double value = 1;
}

message ScalarRange {
  ScalarLimit min = 1;
  ScalarLimit max = 2;
}

message Tuple {
  string name = 1;
  repeated Space value = 2;
}

message Scalar {
  string name = 1;
  oneof range {
    ScalarRange int64_range = 2;
    ScalarRange double_range = 3;
  }
}

message Discrete {
  string name = 1;
  int64 n = 2;
}

message ScalarSequence {
  string name = 1;
  ScalarRange length_range = 2;
  ScalarRange scalar_range = 3;
}

message NamedDiscrete {
  string name = 1;
  repeated string items = 2;
}

message Space {
  string name = 1;
  oneof value {
    Space action_space = 2;
    Tuple tuple = 3;
    Scalar scalar = 4;
    Discrete discrete = 5;
    ScalarSequence scalar_sequence = 6;
    NamedDiscrete named_discrete = 7;
  }
}

I don't know how open you are to breaking changes at this stage.

[ChrisCummins]

I don't know how open you are to breaking changes at this stage.

I am open to breaking changes. The proto interface is described as somewhat stable 😉 . For reference, the last breaking change was in September (#369). It's not too hard to update the in-tree environment implementations and won't affect any user-facing APIs.

Some initial thoughts on your proposed schema:

• I would split the union type Scalar into Int64Scalar and DoubleScalar so as to avoid the need for an extra oneof type check. While we're at it, it would probably be best to replace ScalarRange and ScalarLimit with equivalent Int64Range, DoubleRange etc classes. The current approach of treating all limits as doubles is a bit hacky.
• What is the purpose of Space.action_space? It seems like a weaker version of Space.tuple.
• Tuple and similar classes should be named TupleSpace to distinguish them from the payload messages.
• Do you want to add a DictSpace?
• I think only a single name field is required for each observation and action space. How about dropping all of the name fields from the messages, and keeping the ActionSpace and ObservationSpace wrappers to store the name, and any other relevant fields:

message Space {
  ...
}

message ActionSpace {
  string name = 1;
  Space space = 2;
}

message ObservationSpace {
  string name = 1;
  Space space = 2;
  bool deterministic = 3;
  bool platform_dependent = 4;
  Observation default_value = 5;
}

Overall I think your idea of making the proto schema more closely resemble the Python classes is great. The current schema has evolved over time and has become a bit bloated, there's definitely room for improvement.

Happy holidays! Apologies I will be slow to respond to messages until mid January.

Cheers,
Chris

[sogartar]

I will make a PR that introduces the new message structure and refactor the existing compiler environments. First I will propose the change to the protobuf message descriptions. Once the details are settled, I will refactor the environments.

• Tuple and similar classes should be named TupleSpace to distinguish them from the payload messages.

I thought of putting them in a separate package like it is with the spaces interface in CompilerGym and OpenAI Gym. For example compiler_gym.service.proto.spaces. I don't know how well protobuf bindings work with package nesting in different languages. It seems there is no support for optional python_package.

• I think only a single name field is required for each observation and action space. How about dropping all of the name fields from the messages, and keeping the ActionSpace and ObservationSpace wrappers to store the name, and any other relevant fields:

Nested Spaces should also support names according to Compiler Gym's extension of OpenAI Gym's interface. It can be removed in all other places, but should be kept in Space.

is_commandline in NamedDiscreteSpace is information that is very specific to the particular environment.

// A discrete space in which every point in the space is named. This can be used
// for representing enumerations.
message NamedDiscreteSpace {
  // A list of names for every value in the space. The order in which these
  // values are returned is used.
  repeated string value = 1;
  // If true, the set of space can be interpreted as a list of command line
  // flags.
  bool is_commandline = 2;
}

The interface should allow to represent such information without modification of the gRPC interface.
A similar argument could be made for deterministic and platform_dependent in ObservationSpace. Although, the case to have them there is strong since this information is relevant for every environment.

The Core Python API is considered stable, but it depends on OpenAI Gym, which is in alpha and I think a few months back they introduced breaking changes.
@ChrisCummins, how open are you to non-breaking extensions there?
Fundamentally, action and observation spaces can be very complex and representing them would always be a challenge.
On the environment I am working on the action space contains a permutation. I want this to be somehow explicit in the action space so that RL algorithms and action space wrappers can interpret this in a general way. RL algorithms should be plugable as is the intention of OpenAI Gym.

[sogartar]

When doing the refactoring of the protobuf messages I thought more about it and it seems that there could always be the need to extend the messages. Normal protobuf messages are statically typed, so extensions would always require a modification in CompilerGym. I saw this article that deals with this problem. Basically, you can use Struct or Any in the message where you want to have an arbitrary structure. Then on both sides of the wire this structure is interpreted accordingly.
The python interface of Compiler Gym does not suffer from this problem, because you can always add another class that inherits Space without changing anything else.
@ChrisCummins, do you think this would be too abstract?
At least we can leave the door open by adding such Struct or Any fields in some of the oneof lists. For example in spaces.

[ChrisCummins]

it seems that there could always be the need to extend the messages.

I agree, but just a small clarification: the aim of the proto schema is to provide a flexible enough set of messages that you can embed new data types into, not extend the schema to add new types. You shouldn't have to modify the proto schema to add support for a new data type.

Normal protobuf messages are statically typed, so extensions would always require a modification in CompilerGym... you can use Struct or Any in the message where you want to have an arbitrary structure. Then on both sides of the wire this structure is interpreted accordingly.

There is already a mechanism to do this in CompilerGym, albeit only for observations, not actions. The ObservationSpace message has a opaque_data_format field which, if set, can be used to identify a data type embedded in an Observation. The idea is to use the string_value or binary_value fields to embed a serialized representation of the data:

CompilerGym/compiler_gym/service/proto/compiler_gym_service.proto

Lines 250 to 256 in 9734103

	// An optional string describing an opaque data format, e.g. a data structure
	// that is serialized to a string/binary array for transmission back to the
	// client. It is up to the client and service to agree on how to decode
	// observations using this value. For example, an opaque_data_format of
	// "string_json" could be used to indicate that the observation is a
	// string-serialized JSON value.
	string opaque_data_format = 6;

The ObservationSpaceSpec.from_proto() method contains the logic for how to convert protos to equivalent python types is here. For example, here is the deserialization logic for embedded networkx graphs:

CompilerGym/compiler_gym/views/observation_space_spec.py

Lines 143 to 162 in 9734103

	# Translate from protocol buffer specification to python. There are
	# three variables to derive:
	# (1) space: the gym.Space instance describing the space.
	# (2) translate: is a callback that translates from an Observation
	# message to a python type.
	# (3) to_string: is a callback that translates from a python type to a
	# string for printing.
	if proto.opaque_data_format == "json://networkx/MultiDiGraph":
	# TODO(cummins): Add a Graph space.
	space = make_seq(proto.string_size_range, str, (0, None))
	def translate(observation):
	return nx.readwrite.json_graph.node_link_graph(
	json.loads(observation.string_value), multigraph=True, directed=True
	)
	def to_string(observation):
	return json.dumps(
	nx.readwrite.json_graph.node_link_data(observation), indent=2
	)

For an example of how this is used, the LLVM environment declares the networkx graph space here and serializes and embeds the graphs here.

There are two downsides to this:

1. Adding support for a new opaque data type requires ObservationSpaceSpec.from_proto() which is part of the core library. This means people need to fork+modify CompilerGym to add new opaque data types. We need to add a mechanism by which new deserialization logic can be added without having to modify the base library.
2. The method of serializing to string/bytes may not be as efficient as the Struct / Any types you referenced. In fact, in the case of the embedded networkx graphs above, the data structure starts of as a protobuf from a different schema, which is then converted to JSON, and finally deserialized as networkx. Perhaps adding Struct/Any support could simplify those conversion steps.

So with all that said, I would propose you make the changes in this order:

1. Merging space definitions for action/observations schema as per your suggestions.
2. Adding a mechanism for environments to declare new embedded types, without having to modify ObservationSpaceSpec.from_proto().
3. Investigating whether Struct/Any would provide a performance improvement over embedded strings/bytes, and if so, extending the schema to add this type.

Hopefully this rambly comment has helped give some context as to the design decisions of the current implementation 🙂. I will be out-of-office now until Jan 12th, but happy to discuss further when I'm back, perhaps on a call.

Cheers,
Chris

[sogartar]

2. The method of serializing to string/bytes may not be as efficient as the Struct / Any types you referenced. In fact, in the case of the embedded networkx graphs above, the data structure starts of as a protobuf from a different schema, which is then converted to JSON, and finally deserialized as networkx. Perhaps adding Struct/Any support could simplify those conversion steps.

Any is more wire efficient for larger messages. It is a serialized protobuf message + type URL. Struct is like JSON, so it will embed field names on the wire. In the article that I linked previously, it says

The official language specific protobuf libraries do not have first-class support for converting between Struct and arbitrary protobuf message types, instead it’s always necessary to round-trip via JSON (de)serialization ops.

I don't know if that is still the case as this article is from a few years back.

Any will be convenient if you are using protobuf for the opaque data. If another encoding scheme is used then it will have to be wrapped in a protobuf message. Something like.

message MyOpaqueType {
  bytes data = 1;
}

In the protobuf documentation for Any it says

Currently the runtime libraries for working with Any types are under development.

C++ and Python have support for Any. I don't know what is the situation in other languages.

1. Adding support for a new opaque data type requires ObservationSpaceSpec.from_proto() which is part of the core library. This means people need to fork+modify CompilerGym to add new opaque data types. We need to add a mechanism by which new deserialization logic can be added without having to modify the base library.

I will think about it and present a more concrete approach. It will probably be a mechanism to register deserializers and dispatch based on the type URL from the Any message.

In place of Any we could use something like below.

message Opaque {
  string type_id = 1;
  bytes data = 2;
}

[sogartar]

@ChrisCummins, the Python API is considered stable, but compiler_gym.spaces.Sequence has the problem of str and bytes dtype.
Solving the bytes problem is easy as a byte is just int8.
The string case is a bit more complex because Python does not have character data type and strings can use 1, 2 or 4 bytes per character depending on the encoding. One option is to just use some int type with a given encoding. Then the other side can interpret the data correctly.
Do you think these inconsistencies should be solved?

[ChrisCummins]

Hi @sogartar, yes the inconsistency is that Sequence.dtype should be equivalent to the scalar element type (e.g. type(x[0])), but for bytes sequences it is the type of the entire sequence. Instead it should be the second value below:

In [3]: type(b'abc').__name__
Out[3]: 'bytes'

In [4]: type(b'abc'[0]).__name__
Out[4]: 'int'

Now that I'm looking into it, the string sequence dtype is fine, since Python doesn't distinguish between string / character types, so they can just be strings:

In [1]: type('abc').__name__
Out[1]: 'str'

In [2]: type('abc'[0]).__name__
Out[2]: 'str'

If you're planning on making breaking changes to the proto API then it could also be a good time to clear up that inconsistency between bytes/int dtype for byte arrays.

Cheers,
Chris

[ChrisCummins]

In place of Any we could use something like below.

message Opaque {
  string type_id = 1;
  bytes data = 2;
}

Why not both? :-)

message Event {
  string type_id = 1;
  oneof {
    int32 int32_value = 2;
    ...
    bytes bytes_value = 10;
    Any any_value = 11;
  }
}

Then dispatch to any user-registered deserializer off the type_id value, falling back to the default deserializers.

That means that implementations can use the Any type if appropriate, or they could implement their own custom deserializer using one of the existing field types, e.g. type_id="adjacency_matrix" which could read an adjacency matrix from a list of ints, deserialize some opaque byte array, etc, etc.

Cheers,
Chris

[ChrisCummins]

BTW thanks for your patience, I'm back from vacation now so will be more responsive.

Cheers,
Chris

[sogartar]

@ChrisCummins, using dtype=str in the Sequence space to denote a string will be ambiguous with a sequence of strings. Same holds for bytes. Is it a bytes sequence or a byte sequence?
For a byte sequence it is easy to fix because you can use numpy.int8.
Using int dtype for strings will be wasteful. Then to be correct the implementation would have to handle strings and int collections.

[sogartar]

Why not both? :-)

message Event {
  string type_id = 1;
  oneof {
    int32 int32_value = 2;
    ...
    bytes bytes_value = 10;
    Any any_value = 11;
  }
}

Then dispatch to any user-registered deserializer off the type_id value, falling back to the default deserializers.

OK, I will add the type_id.

[sogartar]

My idea was that if someone wants to have some type like an adjacency matrix they would have in protobuf

message AdjacencyMatrix {
  ...
}

Then use the any_value to transmit the data. Any already carries type information.

[ChrisCummins]

Hi @sogartar, I was browsing through the gym sources today and noticed an interesting use of singledispatch to write a function for working with various different types. Might this be useful for the proto -> python deserializers? Maybe users could register their custom types using a mechanism like this. What do you think?

Cheers,
Chris

[sogartar]

We can use this. It would be more straight forward for someone coming from outside and exploring the code base.
One shortcoming of this approach would be that the dispatch is global and not context dependent. Different environments may have different desires how to translate the data.

[ChrisCummins]

Oh that's a good point. Can different environment class instances bind to a unique dispatcher?

[sogartar]

For observations I think it is already possible to do that by overriding _observation_view_type. It is a bit convoluted though.

In the PR we are preparing for the MLIR environment I added customization of the action and action space conversion to CompilerEnv. The mechanism is a bit different.

The problem boils down to being able to pass/customize 4 functions.

• proto observation space -> gym observation space
• proto action space -> gym action space
• gym action -> proto action
• proto observation -> gym obeservation

I think the best approach is to either do that through inheritance and overriding or passing them to CompilerEnv.__init__.

[sogartar]

The default would be to instantiate the default converters from py_converters, which also support customization by adding additional types and dispatching based on type_id.
Or they could be entirely overridden.